Implement wrap_jax and rename as_op to wrap_py (#1614)

---------

Co-authored-by: Jonas <jonas.dehning@ds.mpg.de>
Co-authored-by: Ricardo Vieira <28983449+ricardoV94@users.noreply.github.com>

Author: 3 people

.github/workflows/test.yml

@@ -208,7 +208,7 @@ jobs:
             micromamba install --yes -q "python~=${PYTHON_VERSION}" mkl "numpy${NUMPY_VERSION}" scipy pip mkl-service graphviz cython pytest coverage pytest-cov pytest-benchmark pytest-mock pytest-sphinx;
           fi
           if [[ $INSTALL_NUMBA == "1" ]]; then micromamba install --yes -q -c conda-forge "python~=${PYTHON_VERSION}" "numba>=0.57"; fi
-          if [[ $INSTALL_JAX == "1" ]]; then micromamba install --yes -q -c conda-forge "python~=${PYTHON_VERSION}" jax jaxlib numpyro && pip install tfp-nightly; fi
+          if [[ $INSTALL_JAX == "1" ]]; then micromamba install --yes -q -c conda-forge "python~=${PYTHON_VERSION}" jax jaxlib numpyro equinox && pip install tfp-nightly; fi
           if [[ $INSTALL_TORCH == "1" ]]; then micromamba install --yes -q -c conda-forge "python~=${PYTHON_VERSION}" pytorch pytorch-cuda=12.1 "mkl<=2024.0" -c pytorch -c nvidia; fi
           if [[ $INSTALL_XARRAY == "1" ]]; then micromamba install --yes -q -c conda-forge "python~=${PYTHON_VERSION}" xarray xarray-einstats; fi
 

doc/environment.yml

@@ -25,4 +25,4 @@ dependencies:
   - ablog
   - pip
   - pip:
-    - -e ..
+    - -e ..[jax]

doc/extending/creating_an_op.rst

@@ -803,10 +803,10 @@ You can omit the :meth:`Rop` functions. Try to implement the testing apparatus d
 :download:`Solution<extending_pytensor_solution_1.py>`
 
 
-:func:`as_op`
+:func:`wrap_py`
 -------------
 
-:func:`as_op` is a Python decorator that converts a Python function into a
+:func:`wrap_py` is a Python decorator that converts a Python function into a
 basic PyTensor :class:`Op` that will call the supplied function during execution.
 
 This isn't the recommended way to build an :class:`Op`, but allows for a quick implementation.
@@ -839,11 +839,11 @@ It takes an optional :meth:`infer_shape` parameter that must have this signature
     inputs PyTensor variables that were declared.
 
 .. note::
-    The python function wrapped by the :func:`as_op` decorator needs to return a new
+    The python function wrapped by the :func:`wrap_py` decorator needs to return a new
     data allocation, no views or in place modification of the input.
 
 
-:func:`as_op` Example
+:func:`wrap_py` Example
 ^^^^^^^^^^^^^^^^^^^^^
 
 .. testcode:: asop
@@ -852,14 +852,14 @@ It takes an optional :meth:`infer_shape` parameter that must have this signature
     import pytensor.tensor as pt
     import numpy as np
     from pytensor import function
-    from pytensor.compile.ops import as_op
+    from pytensor.compile.ops import wrap_py
 
     def infer_shape_numpy_dot(fgraph, node, input_shapes):
         ashp, bshp = input_shapes
         return [ashp[:-1] + bshp[-1:]]
 
 
-    @as_op(
+    @wrap_py(
         itypes=[pt.dmatrix, pt.dmatrix],
         otypes=[pt.dmatrix],
         infer_shape=infer_shape_numpy_dot,

doc/extending/extending_pytensor_solution_1.py

@@ -167,17 +167,17 @@ def test_infer_shape(self):
 
 import numpy as np
 
-# as_op exercice
+# wrap_py exercice
 import pytensor
-from pytensor.compile.ops import as_op
+from pytensor.compile.ops import wrap_py
 
 
 def infer_shape_numpy_dot(fgraph, node, input_shapes):
     ashp, bshp = input_shapes
     return [ashp[:-1] + bshp[-1:]]
 
 
-@as_op(
+@wrap_py(
     itypes=[pt.fmatrix, pt.fmatrix],
     otypes=[pt.fmatrix],
     infer_shape=infer_shape_numpy_dot,
@@ -192,7 +192,7 @@ def infer_shape_numpy_add_sub(fgraph, node, input_shapes):
     return [ashp[0]]
 
 
-@as_op(
+@wrap_py(
     itypes=[pt.fmatrix, pt.fmatrix],
     otypes=[pt.fmatrix],
     infer_shape=infer_shape_numpy_add_sub,
@@ -201,7 +201,7 @@ def numpy_add(a, b):
     return np.add(a, b)
 
 
-@as_op(
+@wrap_py(
     itypes=[pt.fmatrix, pt.fmatrix],
     otypes=[pt.fmatrix],
     infer_shape=infer_shape_numpy_add_sub,

doc/library/index.rst

@@ -61,10 +61,16 @@ Convert to Variable
 
 .. autofunction:: pytensor.as_symbolic(...)
 
+Wrap JAX functions
+==================
+
+.. autofunction:: wrap_jax(...)
+
+   Alias for :func:`pytensor.link.jax.ops.wrap_jax`
+
 Debug
 =====
 
 .. autofunction:: pytensor.dprint(...)
 
    Alias for :func:`pytensor.printing.debugprint`
-

pytensor/__init__.py

@@ -166,7 +166,7 @@ def get_underlying_scalar_constant(v):
 from pytensor.scan.basic import scan
 from pytensor.scan.views import foldl, foldr, map, reduce
 from pytensor.compile.builders import OpFromGraph
-
+from pytensor.link.jax.ops import wrap_jax
 # isort: on
 
 

pytensor/compile/__init__.py

@@ -56,6 +56,7 @@
     register_deep_copy_op_c_code,
     register_view_op_c_code,
     view_op,
+    wrap_py,
 )
 from pytensor.compile.profiling import ProfileStats
 from pytensor.compile.sharedvalue import SharedVariable, shared, shared_constructor

pytensor/compile/ops.py

@@ -1,6 +1,6 @@
 """
 This file contains auxiliary Ops, used during the compilation phase and Ops
-building class (:class:`FromFunctionOp`) and decorator (:func:`as_op`) that
+building class (:class:`FromFunctionOp`) and decorator (:func:`wrap_py`) that
 help make new Ops more rapidly.
 
 """
@@ -268,12 +268,12 @@ def __reduce__(self):
             obj = load_back(mod, name)
         except (ImportError, KeyError, AttributeError):
             raise pickle.PicklingError(
-                f"Can't pickle as_op(), not found as {mod}.{name}"
+                f"Can't pickle wrap_py(), not found as {mod}.{name}"
             )
         else:
             if obj is not self:
                 raise pickle.PicklingError(
-                    f"Can't pickle as_op(), not the object at {mod}.{name}"
+                    f"Can't pickle wrap_py(), not the object at {mod}.{name}"
                 )
         return load_back, (mod, name)
 
@@ -282,6 +282,18 @@ def _infer_shape(self, fgraph, node, input_shapes):
 
 
 def as_op(itypes, otypes, infer_shape=None):
+    import warnings
+
+    warnings.warn(
+        "pytensor.as_op is deprecated and will be removed in a future release. "
+        "Please use pytensor.wrap_py instead.",
+        FutureWarning,
+        stacklevel=2,
+    )
+    return wrap_py(itypes, otypes, infer_shape)
+
+
+def wrap_py(itypes, otypes, infer_shape=None):
     """
     Decorator that converts a function into a basic PyTensor op that will call
     the supplied function as its implementation.
@@ -301,8 +313,8 @@ def infer_shape(fgraph, node, input_shapes):
 
     Examples
     --------
-    @as_op(itypes=[pytensor.tensor.fmatrix, pytensor.tensor.fmatrix],
-           otypes=[pytensor.tensor.fmatrix])
+    @wrap_py(itypes=[pytensor.tensor.fmatrix, pytensor.tensor.fmatrix],
+             otypes=[pytensor.tensor.fmatrix])
     def numpy_dot(a, b):
         return numpy.dot(a, b)
 

pytensor/link/jax/dispatch/basic.py

@@ -13,6 +13,7 @@
 from pytensor.graph import Constant
 from pytensor.graph.fg import FunctionGraph
 from pytensor.ifelse import IfElse
+from pytensor.link.jax.ops import JAXOp
 from pytensor.link.utils import fgraph_to_python
 from pytensor.raise_op import CheckAndRaise
 
@@ -142,3 +143,8 @@ def opfromgraph(*inputs):
             return fgraph_fn(*inputs)
 
     return opfromgraph
+
+
+@jax_funcify.register(JAXOp)
+def jax_op_funcify(op, **kwargs):
+    return op.perform_jax