[pjrt] Do not ignore shape layout in CommonPjRtClient::CreateViewOfDeviceBuffer

superbobry · Google-ML-Automation · commit d0ecd2a9bb98 · 2025-10-07T14:42:29.000-07:00
This allows JAX to support importing DLPack tensors with non-default layouts

PiperOrigin-RevId: 816385564
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -64,7 +64,8 @@ When releasing, please add the new-release-boilerplate to docs/pallas/CHANGELOG.
 * Changes
   * `jax.grad` and `jax.vjp` will now round always primals to float32 if float64
     mode is not enabled.
-
+  * {func}`jax.dlpack.from_dlpack` now accepts arrays with non-default layouts,
+    for example, transposed.
 
 ## JAX 0.7.2 (September 16, 2025)
 
diff --git a/jaxlib/dlpack.cc b/jaxlib/dlpack.cc
@@ -347,22 +347,6 @@ absl::StatusOr<nb::object> DLPackManagedTensorToBuffer(
   }
   xla::Shape shape = xla::ShapeUtil::MakeShapeWithDenseLayout(
       element_type, dimensions, minor_to_major);
-  // Raise an error if the resulting xla::PjRtBuffer would have a non-default
-  // layout.
-  // TODO(skyewm): we do this because JAX doesn't currently have good support
-  // for non-default layouts, and will return wrong results if a non-default
-  // layout is passed to a computation expecting default layouts. Remove this
-  // special case when non-default layouts are better supported by JAX.
-  TF_ASSIGN_OR_RETURN(xla::Layout default_layout,
-                      device->pjrt_device()->client()->GetDefaultLayout(
-                          element_type, dimensions));
-  if (shape.layout() != default_layout) {
-    return xla::Unimplemented(
-        "from_dlpack got array with non-default layout with minor-to-major "
-        "dimensions (%s), expected (%s)",
-        absl::StrJoin(shape.layout().minor_to_major(), ","),
-        absl::StrJoin(default_layout.minor_to_major(), ","));
-  }
 
   TF_ASSIGN_OR_RETURN(auto pjrt_buffer,
                       MakePjrtBuffer(*device->pjrt_device(), dlmt, shape,
diff --git a/tests/array_interoperability_test.py b/tests/array_interoperability_test.py
@@ -15,6 +15,7 @@
 import unittest
 
 from absl.testing import absltest
+import numpy as np
 
 import jax
 import jax.dlpack
@@ -23,8 +24,7 @@
 from jax._src import config
 from jax._src import dlpack as dlpack_src
 from jax._src import test_util as jtu
-
-import numpy as np
+from jax._src.lib import version as jaxlib_version
 
 config.parse_flags_with_absl()
 
@@ -173,6 +173,15 @@ def testTensorFlowToJaxInt64(self):
     dtype_expected = jnp.int64 if config.enable_x64.value else jnp.int32
     self.assertEqual(x.dtype, dtype_expected)
 
+  @unittest.skipIf(not tf, "Test requires TensorFlow")
+  def testTensorFlowToJaxNondefaultLayout(self):
+    if jaxlib_version < (0, 8, 0):
+      self.skipTest(
+          "Non-default layout support requires jaxlib 0.8.0 or newer"
+      )
+    x = tf.transpose(np.arange(4).reshape(2, 2))
+    self.assertAllClose(x.numpy(), jax.dlpack.from_dlpack(x))
+
   @jtu.sample_product(shape=all_shapes, dtype=numpy_dtypes, copy=[False, True])
   def testNumpyToJax(self, shape, dtype, copy):
     rng = jtu.rand_default(self.rng())
@@ -186,28 +195,22 @@ def testNumpyToJax(self, shape, dtype, copy):
     else:
       self.assertAllClose(x_np, _from_dlpack())
 
-  @jtu.sample_product(
-    shape=all_shapes,
-    dtype=numpy_dtypes,
-  )
-  @jtu.run_on_devices("cpu") # NumPy only accepts cpu DLPacks
+  def testNumpyToJaxNondefaultLayout(self):
+    if jaxlib_version < (0, 8, 0):
+      self.skipTest(
+          "Non-default layout support requires jaxlib 0.8.0 or newer"
+      )
+    x = np.arange(4).reshape(2, 2).T
+    self.assertAllClose(x, jax.dlpack.from_dlpack(x))
+
+  @jtu.sample_product(shape=all_shapes, dtype=numpy_dtypes)
+  @jtu.run_on_devices("cpu")  # NumPy only accepts cpu DLPacks
   def testJaxToNumpy(self, shape, dtype):
     rng = jtu.rand_default(self.rng())
     x_jax = jnp.array(rng(shape, dtype))
     x_np = np.from_dlpack(x_jax)
     self.assertAllClose(x_np, x_jax)
 
-  def testNondefaultLayout(self):
-    # Generate numpy array with nonstandard layout
-    a = np.arange(4).reshape(2, 2)
-    b = a.T
-    with self.assertRaisesRegex(
-        RuntimeError,
-        "from_dlpack got array with non-default layout with minor-to-major "
-        r"dimensions \(0,1\), expected \(1,0\)",
-    ):
-      jax.dlpack.from_dlpack(b)
-
 
 class CudaArrayInterfaceTest(jtu.JaxTestCase):
 
diff --git a/tests/pytorch_interoperability_test.py b/tests/pytorch_interoperability_test.py
@@ -17,11 +17,12 @@
 from absl.testing import absltest
 
 import jax
-import jax.dlpack
 from jax._src import config
 from jax._src import test_util as jtu
 from jax._src import xla_bridge
+from jax._src.lib import version as jaxlib_version
 from jax._src.lib import xla_client
+import jax.dlpack
 import jax.numpy as jnp
 
 config.parse_flags_with_absl()
@@ -114,6 +115,15 @@ def testTorchToJaxInt64(self):
     dtype_expected = jnp.int64 if config.enable_x64.value else jnp.int32
     self.assertEqual(x.dtype, dtype_expected)
 
+  def testTorchToJaxNondefaultLayout(self):
+    if jaxlib_version < (0, 8, 0):
+      self.skipTest(
+          "Non-default layout support requires jaxlib 0.8.0 or newer"
+      )
+    x = torch.arange(4).reshape(2, 2).T
+    x = x.cuda() if jtu.test_device_matches(["gpu"]) else x
+    self.assertAllClose(x.cpu().numpy(), jax.dlpack.from_dlpack(x))
+
   @jtu.sample_product(shape=all_shapes, dtype=torch_dtypes)
   def testTorchToJax(self, shape, dtype):
     if not config.enable_x64.value and dtype in [
@@ -130,7 +140,6 @@ def testTorchToJax(self, shape, dtype):
     else:
       x = torch.tensor(x_np)
     x = x.cuda() if jtu.test_device_matches(["gpu"]) else x
-    x = x.contiguous()
     y = jax.dlpack.from_dlpack(x)
     self.assertAllClose(x_np, y)
 
@@ -154,7 +163,6 @@ def testTorchToJaxArray(self, shape, dtype):
     else:
       x = torch.tensor(x_np)
     x = x.cuda() if jtu.test_device_matches(["gpu"]) else x
-    x = x.contiguous()
     y = jax.dlpack.from_dlpack(x)
     self.assertAllClose(x_np, y)
 
