Allow 64-bit output types from ffi_call regardless of enable_x64 flag.

Author: dfm

jax/_src/extend/ffi.py

@@ -27,7 +27,7 @@
 from jax._src import dispatch
 from jax._src import effects
 from jax._src import util
-from jax._src.callback import _check_shape_dtype, callback_batching_rule
+from jax._src.callback import callback_batching_rule
 from jax._src.interpreters import ad
 from jax._src.interpreters import batching
 from jax._src.interpreters import mlir
@@ -209,11 +209,14 @@ def _lowering(
 
 def _result_avals(results: Sequence[ResultMetadata]) -> tuple[core.AbstractValue, ...]:
   avals: list[core.AbstractValue] = []
-  for result in results:
+  for idx, result in enumerate(results):
     if isinstance(result, core.AbstractToken):
       avals.append(result)
     else:
-      _check_shape_dtype(result)
+      if not hasattr(result, "shape") or not hasattr(result, "dtype"):
+        raise ValueError(
+            "All elements of result_shape_dtypes must have 'shape' and 'dtype' "
+            f"attributes. Got {result} at position {idx}.")
       avals.append(core.ShapedArray(result.shape, result.dtype))
   return tuple(avals)
 

tests/extend_test.py

@@ -27,6 +27,7 @@
 
 from jax._src import abstract_arrays
 from jax._src import api
+from jax._src import config
 from jax._src import core
 from jax._src import linear_util
 from jax._src import prng
@@ -326,6 +327,21 @@ def fun(x):
                                 "The use of ffi_call attributes requires"):
       jax.jit(fun).lower(jnp.ones(5)).as_text()
 
+  def testAllow64(self):
+    if config.enable_x64.value:
+      self.skipTest("Requires enable_x64=False")
+    def fun():
+      return jex.ffi.ffi_call("test", jax.ShapeDtypeStruct((), np.int64))()
+    self.assertIn("tensor<i64>", jax.jit(fun).lower().as_text())
+
+  def testInvalidResultType(self):
+    with self.assertRaisesRegex(
+        ValueError, "All elements of result_shape_dtypes.*position 0"):
+      jex.ffi.ffi_call("test", None)()
+    with self.assertRaisesRegex(
+        ValueError, "All elements of result_shape_dtypes.*position 1"):
+      jex.ffi.ffi_call("test", (jax.ShapeDtypeStruct((), np.float32), ()))()
+
 
 def ffi_call_geqrf(x, **kwargs):
   if jtu.test_device_matches(["cpu"]):