Update semantics of to_dlpack. (#2707)

to_dlpack now takes ownership of the original buffer, leaving it in an invalid state.

Author: hawkinsp

jax/dlpack.py

@@ -21,7 +21,8 @@
 def to_dlpack(x: xla.DeviceArray):
   """Returns a DLPack tensor that encapsulates a DeviceArray `x`.
 
-  The DLPack shares memory with `x`.
+  Takes ownership of the contents of `x`; leaves `x` in an invalid/deleted
+  state.
 
   Args:
     x: a `DeviceArray`, on either CPU or GPU.

tests/array_interoperability_test.py

@@ -69,7 +69,7 @@ def testJaxRoundTrip(self, shape, dtype):
     x = jnp.array(np)
     dlpack = jax.dlpack.to_dlpack(x)
     y = jax.dlpack.from_dlpack(dlpack)
-    self.assertAllClose(x, y, check_dtypes=True)
+    self.assertAllClose(np, y, check_dtypes=True)
 
     self.assertRaisesRegex(RuntimeError,
                            "DLPack tensor may be consumed at most once",