pytensor 2.30.2__tar.gz → 2.30.3__tar.gz

{pytensor-2.30.2/pytensor.egg-info → pytensor-2.30.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pytensor
-Version: 2.30.2
+Version: 2.30.3
 Summary: Optimizing compiler for evaluating mathematical expressions on CPUs and GPUs.
 Author-email: pymc-devs <pymc.devs@gmail.com>
 License: .. _license:

{pytensor-2.30.2 → pytensor-2.30.3}/pytensor/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2025-04-01T22:32:44+0200",
+ "date": "2025-04-09T15:18:32+0200",
  "dirty": false,
  "error": null,
- "full-revisionid": "0f5da80c16d5d053f39204d1a728e70e194bef6c",
- "version": "2.30.2"
+ "full-revisionid": "0c398e345af91d28c8dee7075d6d10cafe404025",
+ "version": "2.30.3"
 }
 '''  # END VERSION_JSON
 

{pytensor-2.30.2 → pytensor-2.30.3}/pytensor/link/numba/dispatch/basic.py

@@ -38,6 +38,7 @@ from pytensor.tensor.blas import BatchedDot
 from pytensor.tensor.math import Dot
 from pytensor.tensor.shape import Reshape, Shape, Shape_i, SpecifyShape
 from pytensor.tensor.slinalg import Solve
+from pytensor.tensor.sort import ArgSortOp, SortOp
 from pytensor.tensor.type import TensorType
 from pytensor.tensor.type_other import MakeSlice, NoneConst
 
@@ -433,6 +434,68 @@ def numba_funcify_Shape_i(op, **kwargs):
     return shape_i
 
 
+@numba_funcify.register(SortOp)
+def numba_funcify_SortOp(op, node, **kwargs):
+    @numba_njit
+    def sort_f(a, axis):
+        axis = axis.item()
+
+        a_swapped = np.swapaxes(a, axis, -1)
+        a_sorted = np.sort(a_swapped)
+        a_sorted_swapped = np.swapaxes(a_sorted, -1, axis)
+
+        return a_sorted_swapped
+
+    if op.kind != "quicksort":
+        warnings.warn(
+            (
+                f'Numba function sort doesn\'t support kind="{op.kind}"'
+                " switching to `quicksort`."
+            ),
+            UserWarning,
+        )
+
+    return sort_f
+
+
+@numba_funcify.register(ArgSortOp)
+def numba_funcify_ArgSortOp(op, node, **kwargs):
+    def argsort_f_kind(kind):
+        @numba_njit
+        def argort_vec(X, axis):
+            axis = axis.item()
+
+            Y = np.swapaxes(X, axis, 0)
+            result = np.empty_like(Y)
+
+            indices = list(np.ndindex(Y.shape[1:]))
+
+            for idx in indices:
+                result[(slice(None), *idx)] = np.argsort(
+                    Y[(slice(None), *idx)], kind=kind
+                )
+
+            result = np.swapaxes(result, 0, axis)
+
+            return result
+
+        return argort_vec
+
+    kind = op.kind
+
+    if kind not in ["quicksort", "mergesort"]:
+        kind = "quicksort"
+        warnings.warn(
+            (
+                f'Numba function argsort doesn\'t support kind="{op.kind}"'
+                " switching to `quicksort`."
+            ),
+            UserWarning,
+        )
+
+    return argsort_f_kind(kind)
+
+
 @numba.extending.intrinsic
 def direct_cast(typingctx, val, typ):
     if isinstance(typ, numba.types.TypeRef):

{pytensor-2.30.2 → pytensor-2.30.3}/pytensor/link/numba/dispatch/tensor_basic.py

@@ -136,10 +136,7 @@ def numba_funcify_Join(op, **kwargs):
 def numba_funcify_Split(op, **kwargs):
     @numba_basic.numba_njit
     def split(tensor, axis, indices):
-        # Work around for https://github.com/numba/numba/issues/8257
-        axis = axis % tensor.ndim
-        axis = numba_basic.to_scalar(axis)
-        return np.split(tensor, np.cumsum(indices)[:-1], axis=axis)
+        return np.split(tensor, np.cumsum(indices)[:-1], axis=axis.item())
 
     return split
 

{pytensor-2.30.2 → pytensor-2.30.3}/pytensor/link/numba/dispatch/vectorize_codegen.py

@@ -265,7 +265,7 @@ def _vectorized(
             ctx.nrt.incref(
                 builder,
                 sig.return_type.types[inplace_idx],
-                outputs[inplace_idx]._get_value(),
+                outputs[inplace_idx]._getvalue(),
             )
         return ctx.make_tuple(
             builder, sig.return_type, [out._getvalue() for out in outputs]

{pytensor-2.30.2 → pytensor-2.30.3}/pytensor/scalar/basic.py

@@ -1302,19 +1302,7 @@ class ScalarOp(COp):
     def __str__(self):
         if hasattr(self, "name") and self.name:
             return self.name
-        else:
-            param = [
-                (k, v)
-                for k, v in self.__dict__.items()
-                if k
-                not in ("name", "_op_use_c_code", "bool", "output_types_preference")
-            ]
-            if param:
-                classname = self.__class__.__name__
-                args = ", ".join(f"{k}={v}" for k, v in param)
-                return f"{classname}{{{args}}}"
-            else:
-                return self.__class__.__name__
+        return self.__class__.__name__
 
     def c_code_cache_version(self):
         return (4,)
@@ -4102,6 +4090,7 @@ class ScalarInnerGraphOp(ScalarOp, HasInnerGraph):
 
     def __init__(self, *args, **kwargs):
         self.prepare_node_called = set()
+        super().__init__(*args, **kwargs)
 
     def _cleanup_graph(self, inputs, outputs):
         # TODO: We could convert to TensorVariable, optimize graph,
@@ -4344,9 +4333,9 @@ class Composite(ScalarInnerGraphOp):
 
         # Rename internal variables
         for i, r in enumerate(self.fgraph.inputs):
-            r.name = f"i{int(i)}"
+            r.name = f"i{i}"
         for i, r in enumerate(self.fgraph.outputs):
-            r.name = f"o{int(i)}"
+            r.name = f"o{i}"
         io = set(self.fgraph.inputs + self.fgraph.outputs)
         for i, r in enumerate(self.fgraph.variables):
             if (
@@ -4354,7 +4343,7 @@ class Composite(ScalarInnerGraphOp):
                 and r not in io
                 and len(self.fgraph.clients[r]) > 1
             ):
-                r.name = f"t{int(i)}"
+                r.name = f"t{i}"
 
         if len(self.fgraph.outputs) > 1 or len(self.fgraph.apply_nodes) > 10:
             self._name = "Composite{...}"
@@ -4441,16 +4430,12 @@ class Composite(ScalarInnerGraphOp):
         if hasattr(self, "_c_code"):
             return self._c_code
 
-        subd = dict(
-            chain(
-                ((e, f"%(i{int(i)})s") for i, e in enumerate(self.fgraph.inputs)),
-                ((e, f"%(o{int(i)})s") for i, e in enumerate(self.fgraph.outputs)),
-            )
-        )
+        fg = self.fgraph
+        subd = {e: f"%(i{i})s" for i, e in enumerate(fg.inputs)}
 
-        for var in self.fgraph.variables:
+        for var in fg.variables:
             if var.owner is None:
-                if var not in self.fgraph.inputs:
+                if var not in fg.inputs:
                     # This is an orphan
                     if isinstance(var, Constant) and isinstance(var.type, CLinkerType):
                         subd[var] = f"({var.type.c_literal(var.data)})"
@@ -4465,30 +4450,35 @@ class Composite(ScalarInnerGraphOp):
                 # flag for elemwise ops to check.
                 self.inner_float16 = True
 
-        _c_code = "{\n"
-        self.nodenames = [
-            f"%(nodename)s_subnode{int(j)}"
-            for j, n in enumerate(self.fgraph.toposort())
-        ]
+        self.nodenames = nodenames = []  # Used by self.c_support_code_apply
 
+        _c_code = "{\n"
         i = 0
-        for j, node in enumerate(self.fgraph.toposort()):
+        for j, node in enumerate(fg.toposort()):
             for output in node.outputs:
                 if output not in subd:
                     i += 1
-                    name = f"V%(id)s_tmp{int(i)}"
+                    name = f"V%(id)s_tmp{i}"
                     subd[output] = name
                     _c_code += f"{output.type.dtype_specs()[1]} {name};\n"
+
+            nodename = f"%(nodename)s_subnode{j}"
+            nodenames.append(nodename)
+
             s = node.op.c_code(
                 node,
-                self.nodenames[j],
+                nodename,
                 [subd[input] for input in node.inputs],
                 [subd[output] for output in node.outputs],
-                dict(fail="%(fail)s", id=f"%(id)s_{int(j)}"),
+                dict(fail="%(fail)s", id=f"%(id)s_{j}"),
             )
             _c_code += s
             _c_code += "\n"
 
+        # Copy the temporary outputs to the real outputs
+        for i, output in enumerate(fg.outputs):
+            _c_code += f"%(o{i})s = {subd[output]};\n"
+
         _c_code += "}\n"
 
         self._c_code = _c_code
@@ -4498,8 +4488,8 @@ class Composite(ScalarInnerGraphOp):
     def c_code(self, node, nodename, inames, onames, sub):
         d = dict(
             chain(
-                zip((f"i{int(i)}" for i in range(len(inames))), inames, strict=True),
-                zip((f"o{int(i)}" for i in range(len(onames))), onames, strict=True),
+                zip((f"i{i}" for i in range(len(inames))), inames, strict=True),
+                zip((f"o{i}" for i in range(len(onames))), onames, strict=True),
             ),
             **sub,
         )
@@ -4512,7 +4502,7 @@ class Composite(ScalarInnerGraphOp):
         return self.c_code_template % d
 
     def c_code_cache_version_outer(self) -> tuple[int, ...]:
-        return (5,)
+        return (6,)
 
 
 class Compositef32:

{pytensor-2.30.2 → pytensor-2.30.3}/pytensor/scalar/loop.py

@@ -55,6 +55,7 @@ class ScalarLoop(ScalarInnerGraphOp):
         constant: Sequence[Variable] | None = None,
         until: Variable | None = None,
         name="ScalarLoop",
+        **kwargs,
     ):
         if constant is None:
             constant = []
@@ -75,7 +76,7 @@ class ScalarLoop(ScalarInnerGraphOp):
         self.nout = len(self.outputs)
         self.name = name
 
-        super().__init__()
+        super().__init__(**kwargs)
 
     def output_types(self, input_types):
         return self.outputs_type
@@ -115,7 +116,7 @@ class ScalarLoop(ScalarInnerGraphOp):
         self._fgraph = fgraph
         return self._fgraph
 
-    def clone(self):
+    def clone(self, name=None, **kwargs):
         if self.is_while:
             *update, until = self.outputs
         else:
@@ -127,7 +128,8 @@ class ScalarLoop(ScalarInnerGraphOp):
             update=update,
             constant=constant,
             until=until,
-            name=self.name,
+            name=self.name if name is None else name,
+            **kwargs,
         )
 
     @property
@@ -135,20 +137,7 @@ class ScalarLoop(ScalarInnerGraphOp):
         raise NotImplementedError
 
     def make_new_inplace(self, output_types_preference=None, name=None):
-        """
-        This op.__init__ fct don't have the same parameter as other scalar op.
-        This break the insert_inplace_optimizer optimization.
-        This fct allow fix patch this.
-
-        """
-        d = {k: getattr(self, k) for k in self.init_param}
-        out = self.__class__(**d)
-        if name:
-            out.name = name
-        else:
-            name = out.name
-        super(ScalarLoop, out).__init__(output_types_preference, name)
-        return out
+        return self.clone(output_types_preference=output_types_preference, name=name)
 
     def make_node(self, n_steps, *inputs):
         assert len(inputs) == self.nin - 1
@@ -223,17 +212,15 @@ class ScalarLoop(ScalarInnerGraphOp):
         # The first input is `n_steps` so we skip it in the mapping dictionary
         n_update = len(self.outputs) - (1 if self.is_while else 0)
         carry_subd = {
-            c: f"%(i{int(i)})s" for i, c in enumerate(fgraph.inputs[:n_update], start=1)
+            c: f"%(i{i})s" for i, c in enumerate(fgraph.inputs[:n_update], start=1)
         }
         constant_subd = {
-            c: f"%(i{int(i)})s"
+            c: f"%(i{i})s"
             for i, c in enumerate(fgraph.inputs[n_update:], start=n_update + 1)
         }
-        update_subd = {
-            u: f"%(o{int(i)})s" for i, u in enumerate(fgraph.outputs[:n_update])
-        }
+        out_subd = {u: f"%(o{i})s" for i, u in enumerate(fgraph.outputs[:n_update])}
         until_subd = {u: "until" for u in fgraph.outputs[n_update:]}
-        subd = {**carry_subd, **constant_subd, **update_subd, **until_subd}
+        subd = {**carry_subd, **constant_subd, **until_subd}
 
         for var in fgraph.variables:
             if var.owner is None:
@@ -257,11 +244,11 @@ class ScalarLoop(ScalarInnerGraphOp):
             _c_code += "bool until = 1;\n\n"
 
         # Copy carried inputs
-        for i, (var, name) in enumerate(carry_subd.items()):
-            copy_var_name = f"{name}_copy{i}"
-            _c_code += f"{var.type.dtype_specs()[1]} {copy_var_name} = {name};\n"
-            carry_subd[var] = copy_var_name
-            subd[var] = copy_var_name
+        for i, (var, name) in enumerate(carry_subd.items(), start=1):
+            carry_var_name = f"{name}_carry{i}"
+            _c_code += f"{var.type.dtype_specs()[1]} {carry_var_name} = {name};\n"
+            carry_subd[var] = carry_var_name
+            subd[var] = carry_var_name
 
         # _c_code += 'printf("inputs=[");'
         # for i in range(1, len(fgraph.inputs)):
@@ -270,34 +257,39 @@ class ScalarLoop(ScalarInnerGraphOp):
 
         _c_code += "\nfor(%(n_steps_dtype)s i = 0; i < %(n_steps)s; i++){\n"
 
-        self.nodenames = [
-            f"%(nodename)s_subnode{int(j)}" for j, n in enumerate(fgraph.toposort())
-        ]
+        # Used by self.c_support_code_apply
+        self.nodenames = nodenames = []
 
         i = 0
         for j, node in enumerate(fgraph.toposort()):
             for output in node.outputs:
                 if output not in subd:
                     i += 1
-                    name = f"V%(id)s_tmp{int(i)}"
+                    name = f"V%(id)s_tmp{i}"
                     subd[output] = name
                     _c_code += f"{output.type.dtype_specs()[1]} {name};\n"
+
+            nodename = f"%(nodename)s_subnode{j}"
+            nodenames.append(nodename)
+
             s = node.op.c_code(
                 node,
-                self.nodenames[j],
+                nodename,
                 # Any node that depended on `init` will depend on `update` instead
                 # The initial value of `update` was set to `init` before the loop
                 [subd[input] for input in node.inputs],
                 [subd[output] for output in node.outputs],
-                dict(fail="%(fail)s", id=f"%(id)s_{int(j)}"),
+                dict(fail="%(fail)s", id=f"%(id)s_{j}"),
             )
             _c_code += s
             _c_code += "\n"
 
-        # Set the carry variables to the output variables
+        # Update the carry variables to the output variables
         _c_code += "\n"
-        for init, update in zip(carry_subd.values(), update_subd.values(), strict=True):
-            _c_code += f"{init} = {update};\n"
+        for carry, out in zip(
+            carry_subd.values(), fgraph.outputs[:n_update], strict=True
+        ):
+            _c_code += f"{carry} = {subd[out]};\n"
 
         # _c_code += 'printf("%%ld\\n", i);\n'
         # for carry in range(1, 10):
@@ -309,6 +301,10 @@ class ScalarLoop(ScalarInnerGraphOp):
         # End of the loop
         _c_code += "}\n"
 
+        # Assign the carry variables to the outputs
+        for out, carry in zip(out_subd.values(), carry_subd.values(), strict=True):
+            _c_code += f"{out} = {carry};\n"
+
         # Output until flag
         if self.is_while:
             _c_code += f"%(o{len(fgraph.outputs)-1})s = until;\n"
@@ -322,8 +318,8 @@ class ScalarLoop(ScalarInnerGraphOp):
     def c_code(self, node, nodename, inames, onames, sub):
         d = dict(
             chain(
-                zip((f"i{int(i)}" for i in range(len(inames))), inames, strict=True),
-                zip((f"o{int(i)}" for i in range(len(onames))), onames, strict=True),
+                zip((f"i{i}" for i in range(len(inames))), inames, strict=True),
+                zip((f"o{i}" for i in range(len(onames))), onames, strict=True),
             ),
             **sub,
         )
@@ -343,4 +339,4 @@ class ScalarLoop(ScalarInnerGraphOp):
         return res
 
     def c_code_cache_version_outer(self):
-        return (3,)
+        return (4,)