pytensor 2.33.0__tar.gz → 2.34.0__tar.gz

{pytensor-2.33.0/pytensor.egg-info → pytensor-2.34.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pytensor
-Version: 2.33.0
+Version: 2.34.0
 Summary: Optimizing compiler for evaluating mathematical expressions on CPUs and GPUs.
 Author-email: pymc-devs <pymc.devs@gmail.com>
 License-Expression: BSD-3-Clause

{pytensor-2.33.0 → pytensor-2.34.0}/doc/environment.yml

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

{pytensor-2.33.0 → pytensor-2.34.0}/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,

{pytensor-2.33.0 → pytensor-2.34.0}/doc/extending/extending_pytensor_solution_1.py

@@ -167,9 +167,9 @@ class TestSumDiffOp(utt.InferShapeTester):
 
 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):
@@ -177,7 +177,7 @@ def infer_shape_numpy_dot(fgraph, node, 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,

{pytensor-2.33.0 → pytensor-2.34.0}/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-2.33.0 → pytensor-2.34.0}/pytensor/__init__.py

@@ -166,7 +166,7 @@ from pytensor.scan import checkpoints
 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-2.33.0 → pytensor-2.34.0}/pytensor/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2025-10-01T04:44:52+0200",
+ "date": "2025-10-07T21:55:32+0200",
  "dirty": false,
  "error": null,
- "full-revisionid": "7779b07b892118151f96d1ed98bd69d4281c687c",
- "version": "2.33.0"
+ "full-revisionid": "b6a9b68e9d9bc14e0b5d095bcbd2756895fe221d",
+ "version": "2.34.0"
 }
 '''  # END VERSION_JSON
 

{pytensor-2.33.0 → pytensor-2.34.0}/pytensor/compile/__init__.py

@@ -56,6 +56,7 @@ from pytensor.compile.ops import (
     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-2.33.0 → pytensor-2.34.0}/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 @@ class FromFunctionOp(Op):
             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 @@ class FromFunctionOp(Op):
 
 
 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 as_op(itypes, otypes, infer_shape=None):
 
     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-2.33.0 → pytensor-2.34.0}/pytensor/link/jax/dispatch/basic.py

@@ -13,6 +13,7 @@ from pytensor.configdefaults import config
 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 jax_funcify_OpFromGraph(ofg: OpFromGraph, node=None, **kwargs) -> Callable:
             return fgraph_fn(*inputs)
 
     return opfromgraph
+
+
+@jax_funcify.register(JAXOp)
+def jax_op_funcify(op, **kwargs):
+    return op.perform_jax