Fix support for custom JAX operations

Author: wesselb

lab/custom.py

@@ -180,18 +180,24 @@ def bvn_cdf(a, b, c):
     # does not work for `bvn_cdf_`. Moreover, we need to ensure that the function
     # runs on `float64s`.
     res_dtype = reduce(np.promote_types, [x.dtype for x in (a, b, c)])
-    res = bvn_cdf_(a.astype(np.float64), b.astype(np.float64), c.astype(np.float64))
+    # The C interface requires NumPy objects of the right data type.
+    res = bvn_cdf_(
+        np.asarray(a).astype(np.float64),
+        np.asarray(b).astype(np.float64),
+        np.asarray(c).astype(np.float64),
+    )
     return res.astype(res_dtype)
 
 
 def s_bvn_cdf(s_y, y, a, b, c):
     res_dtype = reduce(np.promote_types, [x.dtype for x in (s_y, y, a, b, c)])
+    # The C interface requires NumPy objects of the right data type.
     res = s_bvn_cdf_(
-        s_y.astype(np.float64),
-        y.astype(np.float64),
-        a.astype(np.float64),
-        b.astype(np.float64),
-        c.astype(np.float64),
+        np.asarray(s_y).astype(np.float64),
+        np.asarray(y).astype(np.float64),
+        np.asarray(a).astype(np.float64),
+        np.asarray(b).astype(np.float64),
+        np.asarray(c).astype(np.float64),
     )
     return tuple(x.astype(res_dtype) for x in res)
 

lab/jax/custom.py

@@ -1,9 +1,23 @@
 from functools import wraps
 
-import jax.experimental.host_callback as hcb
-from jax import ShapeDtypeStruct, custom_vjp
+from jax import ShapeDtypeStruct
+from jax import __version__ as jax_version
+from jax import custom_vjp
+from packaging.version import Version
 from plum import Dispatcher, convert
 
+if Version(jax_version) >= Version("0.5"):
+    from jax.experimental import io_callback
+
+# IO callbacks do not support JVP in JAX before `0.5`, so we emulate the call with
+# `host_callback`.
+else:  # pragma: no cover
+    from jax.experimental import host_callback
+
+    def io_callback(f, shapes, x):
+        return host_callback.call(f, x, result_shape=shapes)
+
+
 from ..custom import TensorDescription
 
 __all__ = ["jax_register"]
@@ -30,13 +44,13 @@ def parse_inference_result(xs: tuple):
     return tuple(parse_inference_result(x) for x in xs)
 
 
-def _wrap_hcb(f, i_f):
+def _wrap_cb(f, i_f):
     @wraps(f)
     def f_wrapped(*args, **kw_args):
-        return hcb.call(
+        return io_callback(
             lambda x: f(*x[0], **x[1]),
-            arg=(args, kw_args),
-            result_shape=parse_inference_result(i_f(*args, **kw_args)),
+            parse_inference_result(i_f(*args, **kw_args)),
+            (args, kw_args),
         )
 
     return f_wrapped
@@ -55,8 +69,8 @@ def jax_register(f, i_f, s_f, i_s_f):
     Returns:
         function: JAX function.
     """
-    f = _wrap_hcb(f, i_f)
-    s_f = _wrap_hcb(s_f, i_s_f)
+    f = _wrap_cb(f, i_f)
+    s_f = _wrap_cb(s_f, i_s_f)
 
     f = custom_vjp(f)
 

lab/types.py

@@ -194,9 +194,16 @@ def _torch_lookup(dtype):
 
 def _name(x):
     try:
-        return x.name
+        name = x.name
     except AttributeError:
-        return x.__name__
+        name = x.__name__
+
+    # `int`s are `int64`s. We need to do this conversion explicity, because `torch.int`
+    # is `torch.int32` in later versions.
+    if name == "int":
+        name = "int64"
+
+    return name
 
 
 # Add conversions between data types.
@@ -357,7 +364,7 @@ def dtype_int(dtype: DType):
     """
     # TODO: Is there a better way of doing this?
     name = list(convert(dtype, NPDType).__name__)
-    while name and name[0] not in set([str(i) for i in range(10)]):
+    while name and not name[0].isdigit():
         name.pop(0)
     return _convert_back(_name_to_numpy_dtype("int" + "".join(name)), dtype)