add test

Co-authored-by: Adam Paszke <[email protected]>

python/dex/__init__.py

@@ -31,7 +31,7 @@ def _from_ptr(cls, ptr):
     return self
 
   def __del__(self):
-    if api.nofree: return
+    if api is None or api.nofree: return
     api.destroyContext(self)
 
   def __getattr__(self, name):

python/dex/interop/jax2dex.py

@@ -129,6 +129,8 @@ def pprint(self) -> str:
 from jax._src.util import (unzip2, wraps, split_list, partition_list, safe_map,
                            safe_zip, cache)
 
+from .. import eval
+
 map = safe_map
 zip = safe_zip
 
@@ -174,6 +176,7 @@ def dex_fun(*args, **kwargs):
     return tree_unflatten(out_tree, out_flat)
   return dex_fun
 
+# TODO try to delete this, rely on existing jax functions instead
 @cache()
 def make_jaxpr(fun: Callable, in_tree: PyTreeDef,
                in_avals: tuple[core.AbstractValue],  # with DBIdx in them
@@ -238,7 +241,10 @@ def write(v: core.Var, val: str) -> None:
     block = Block([], expr)
   print(jaxpr, end='\n\n')
   print(pprint(expr), end='\n\n')
-  return dex.eval(pprint(expr)).compile()
+  return eval(pprint(expr)).compile()
+
+def pprint(e):
+  return e.pprint()
 
 ExprMaker = Callable[[Any, ...], Expr]
 expr_makers: Dict[core.Primitive, ExprMaker] = {}
@@ -294,3 +300,7 @@ def aval_to_type(aval: core.AbstractValue) -> Type:
   else:
     raise NotImplementedError(aval)
 
+
+###
+
+

python/tests/jax_test.py

@@ -12,6 +12,7 @@
 
 import dex
 from dex.interop.jax import primitive
+from dex.interop.jax2dex import dexjit
 
 
 @unittest.skip
@@ -206,6 +207,16 @@ def grad_jax(x, y):
         jax.jit(grad_dex)(x, y),
         jax.jit(grad_jax)(x, y))
 
+class JAX2DexTest(unittest.TestCase):
+
+  def test_basic(self):
+    @dexjit
+    def f(x, y):
+      assert x.ndim == y.ndim == 0
+      return x + y
+
+    f(1., 2.)
+
 
 if __name__ == "__main__":
   unittest.main()