Move JAX linker to new "link" module (#219)

tests/sandbox/test_jax.py → tests/link/test_jax.py

@@ -52,7 +52,7 @@ def compare_jax_and_py(
         assert_fn = partial(np.testing.assert_allclose, rtol=1e-4)
 
     opts = theano.gof.Query(include=[None], exclude=["cxx_only", "BlasOpt"])
-    jax_mode = theano.compile.mode.Mode(theano.sandbox.jax_linker.JAXLinker(), opts)
+    jax_mode = theano.compile.mode.Mode(theano.link.jax.JAXLinker(), opts)
     py_mode = theano.compile.Mode("py", opts)
 
     theano_jax_fn = theano.function(fgraph.inputs, fgraph.outputs, mode=jax_mode)
@@ -833,7 +833,7 @@ def test_jax_BatchedDot():
     # A dimension mismatch should raise a TypeError for compatibility
     inputs = [get_test_value(a)[:-1], get_test_value(b)]
     opts = theano.gof.Query(include=[None], exclude=["cxx_only", "BlasOpt"])
-    jax_mode = theano.compile.mode.Mode(theano.sandbox.jax_linker.JAXLinker(), opts)
+    jax_mode = theano.compile.mode.Mode(theano.link.jax.JAXLinker(), opts)
     theano_jax_fn = theano.function(fgraph.inputs, fgraph.outputs, mode=jax_mode)
     with pytest.raises(TypeError):
         theano_jax_fn(*inputs)

theano/compile/mode.py

@@ -10,7 +10,7 @@ import theano
 import theano.gof.vm
 from theano import config, gof
 from theano.compile.function.types import Supervisor
-from theano.sandbox.jax_linker import JAXLinker
+from theano.link.jax import JAXLinker
 
 
 _logger = logging.getLogger("theano.compile.mode")

theano/link/jax/__init__.py

+from theano.link.jax.jax_linker import JAXLinker

theano/sandbox/jax_linker.py → theano/link/jax/jax_linker.py

@@ -50,7 +50,7 @@ class JAXLinker(PerformLinker):
         """
         import jax
 
-        from theano.sandbox.jaxify import jax_funcify
+        from theano.link.jax.jax_dispatch import jax_funcify
 
         output_nodes = [o.owner for o in self.fgraph.outputs]