Implement @as_jax_op to wrap a JAX function for use in PyTensor

doc/conf.py

@@ -36,6 +36,7 @@
     "jax": ("https://jax.readthedocs.io/en/latest", None),
     "numpy": ("https://numpy.org/doc/stable", None),
     "torch": ("https://pytorch.org/docs/stable", None),
+    "equinox": ("https://docs.kidger.site/equinox/", None),
 }
 
 needs_sphinx = "3"

doc/environment.yml

@@ -24,4 +24,4 @@ dependencies:
   - pip
   - pip:
     - sphinx_sitemap
-    - -e ..
+    - -e ..[jax]

doc/library/index.rst

@@ -63,6 +63,13 @@ Convert to Variable
 
 .. autofunction:: pytensor.as_symbolic(...)
 
+Wrap JAX functions
+==================
+
+.. autofunction:: as_jax_op(...)
+
+   Alias for :func:`pytensor.link.jax.ops.as_jax_op`
+
 Debug
 =====
 

pyproject.toml

@@ -78,7 +78,7 @@ tests = [
     "pytest-sphinx",
 ]
 rtd = ["sphinx>=5.1.0,<6", "pygments", "pydot", "pydot2", "pydot-ng"]
-jax = ["jax", "jaxlib"]
+jax = ["jax", "jaxlib", "equinox"]
 numba = ["numba>=0.57", "llvmlite"]
 
 [tool.setuptools.packages.find]

pytensor/__init__.py

@@ -167,6 +167,18 @@ def get_underlying_scalar_constant(v):
 from pytensor.scan.views import foldl, foldr, map, reduce
 from pytensor.compile.builders import OpFromGraph
 
+try:
+    import pytensor.link.jax.ops
+    from pytensor.link.jax.ops import as_jax_op
+except ImportError as e:
+    import_error_as_jax_op = e
+
+    def as_jax_op(*args, **kwargs):
+        raise ImportError(
+            "JAX and/or equinox are not installed. Install them"
+            " to use this function: pip install pytensor[jax]"
+        ) from import_error_as_jax_op
+
 # isort: on