pytensor 2.36.1__tar.gz → 2.36.2__tar.gz

{pytensor-2.36.1/pytensor.egg-info → pytensor-2.36.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pytensor
-Version: 2.36.1
+Version: 2.36.2
 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.36.1 → pytensor-2.36.2}/pytensor/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2025-12-21T14:56:52+0000",
+ "date": "2026-01-02T17:20:52+0100",
  "dirty": false,
  "error": null,
- "full-revisionid": "409a4b2a02e64acc0ca12cd0e046e99fd6f535b9",
- "version": "2.36.1"
+ "full-revisionid": "0f864a33958e68e5110d043c4b5ef549cd1d6074",
+ "version": "2.36.2"
 }
 '''  # END VERSION_JSON
 

{pytensor-2.36.1 → pytensor-2.36.2}/pytensor/compile/builders.py

@@ -872,7 +872,7 @@ class OpFromGraph(Op, HasInnerGraph):
 
     def clone(self):
         res = copy(self)
-        res.fgraph = res.fgraph.clone()
+        res.fgraph = res.fgraph.clone(clone_inner_graphs=True)
         return res
 
     def perform(self, node, inputs, outputs):

{pytensor-2.36.1 → pytensor-2.36.2}/pytensor/graph/fg.py

@@ -838,9 +838,13 @@ class FunctionGraph(MetaObject):
     def __repr__(self):
         return f"FunctionGraph({', '.join(graph_as_string(self.inputs, self.outputs))})"
 
-    def clone(self, check_integrity=True) -> "FunctionGraph":
+    def clone(
+        self, check_integrity=True, clone_inner_graphs: bool = False
+    ) -> "FunctionGraph":
         """Clone the graph."""
-        return self.clone_get_equiv(check_integrity)[0]
+        return self.clone_get_equiv(
+            check_integrity, clone_inner_graphs=clone_inner_graphs
+        )[0]
 
     def clone_get_equiv(
         self, check_integrity: bool = True, attach_feature: bool = True, **kwargs

{pytensor-2.36.1 → pytensor-2.36.2}/pytensor/link/numba/dispatch/basic.py

@@ -478,7 +478,7 @@ def numba_funcify_ensure_cache(op, *args, **kwargs) -> tuple[Callable, str | Non
             src=f"def {op_name}(*args): return jitable_func(*args)",
             function_name=op_name,
             global_env=globals() | {"jitable_func": jitable_func},
-            cache_key=cache_key,
+            cache_key=f"{cache_key}_fastmath{int(config.numba__fastmath)}",
         )
         return numba_njit(cached_func, cache=True), cache_key
 

{pytensor-2.36.1 → pytensor-2.36.2}/pytensor/link/numba/dispatch/compile_ops.py

@@ -1,5 +1,6 @@
 from copy import deepcopy
 from hashlib import sha256
+from textwrap import dedent
 
 import numba
 import numpy as np
@@ -10,6 +11,7 @@ from pytensor.compile.io import In, Out
 from pytensor.compile.mode import NUMBA
 from pytensor.compile.ops import DeepCopyOp, TypeCastingOp
 from pytensor.ifelse import IfElse
+from pytensor.link.numba.cache import compile_numba_function_src
 from pytensor.link.numba.dispatch import basic as numba_basic
 from pytensor.link.numba.dispatch.basic import (
     numba_funcify_and_cache_key,
@@ -106,30 +108,35 @@ def numba_funcify_DeepCopyOp(op, node, **kwargs):
 @register_funcify_default_op_cache_key(IfElse)
 def numba_funcify_IfElse(op, **kwargs):
     n_outs = op.n_outs
+    as_view = op.as_view
 
-    if n_outs > 1:
-
-        @numba_basic.numba_njit
-        def ifelse(cond, *args):
-            if cond:
-                res = args[:n_outs]
-            else:
-                res = args[n_outs:]
-
-            return res
+    true_names = [f"t{i}" for i in range(n_outs)]
+    false_names = [f"f{i}" for i in range(n_outs)]
+    arg_list = ", ".join((*true_names, *false_names))
 
+    if as_view:
+        true_returns = ", ".join(true_names)
     else:
+        true_returns = ", ".join(f"{name}.copy()" for name in true_names)
+    # We only ever view (alias) variables from the true branch. False branch variables must always be copied.
+    false_returns = ", ".join(f"{name}.copy()" for name in false_names)
+
+    func_src = dedent(
+        f"""
+            def ifelse(cond, {arg_list}):
+                if cond:
+                    return {true_returns}
+                else:
+                    return {false_returns}
+        """
+    )
 
-        @numba_basic.numba_njit
-        def ifelse(cond, *args):
-            if cond:
-                res = args[:n_outs]
-            else:
-                res = args[n_outs:]
-
-            return res[0]
+    ifelse_func = numba_basic.numba_njit(
+        compile_numba_function_src(func_src, "ifelse", globals())
+    )
 
-    return ifelse
+    cache_version = 1
+    return ifelse_func, cache_version
 
 
 @register_funcify_and_cache_key(CheckAndRaise)

{pytensor-2.36.1 → pytensor-2.36.2}/pytensor/link/numba/dispatch/random.py

@@ -216,14 +216,16 @@ def core_MvNormalRV(op, node):
 
 @numba_core_rv_funcify.register(ptr.DirichletRV)
 def core_DirichletRV(op, node):
+    dtype = op.dtype
+
     @numba_basic.numba_njit
     def random_fn(rng, alpha):
-        y = np.empty_like(alpha)
+        y = np.empty_like(alpha, dtype=dtype)
         for i in range(len(alpha)):
             y[i] = rng.gamma(alpha[i], 1.0)
         return y / y.sum()
 
-    return random_fn
+    return random_fn, 1
 
 
 @numba_core_rv_funcify.register(ptr.GumbelRV)
@@ -410,7 +412,7 @@ def numba_funcify_RandomVariable(op: RandomVariableWithCoreShape, node, **kwargs
     rv_op: RandomVariable = rv_node.op
 
     try:
-        core_rv_fn = numba_core_rv_funcify(rv_op, rv_node)
+        core_rv_fn_and_cache_key = numba_core_rv_funcify(rv_op, rv_node)
     except NotImplementedError:
         py_impl = generate_fallback_impl(rv_op, node=rv_node, **kwargs)
 
@@ -420,6 +422,16 @@ def numba_funcify_RandomVariable(op: RandomVariableWithCoreShape, node, **kwargs
 
         return fallback_rv, None
 
+    match core_rv_fn_and_cache_key:
+        case (core_rv_fn, (int() | None) as core_cache_key):
+            pass
+        case (_core_rv_fn, invalid_core_cache_key):
+            raise ValueError(
+                f"Invalid core_cache_key returned from numba_core_rv_funcify: {invalid_core_cache_key}. Must be int or None."
+            )
+        case core_rv_fn:
+            core_cache_key = "__None__"
+
     size = rv_op.size_param(rv_node)
     dist_params = rv_op.dist_params(rv_node)
     size_len = None if isinstance(size.type, NoneTypeT) else get_vector_length(size)
@@ -469,16 +481,19 @@ def numba_funcify_RandomVariable(op: RandomVariableWithCoreShape, node, **kwargs
 
         return impl
 
-    rv_op_props_dict = rv_op.props_dict() if hasattr(rv_op, "props_dict") else {}
-    random_rv_key_contents = (
-        type(op),
-        type(rv_op),
-        rv_op,
-        tuple(rv_op_props_dict.items()),
-        size_len,
-        core_shape_len,
-        inplace,
-        input_bc_patterns,
-    )
-    random_rv_key = sha256(str(random_rv_key_contents).encode()).hexdigest()
+    if core_cache_key is None:
+        # If the core RV can't be cached, then the whole RV can't be cached
+        random_rv_key = None  # type: ignore[unreachable]
+    else:
+        random_rv_key_contents = (
+            type(op),
+            type(rv_op),
+            tuple(rv_op._props_dict().items()),  # type: ignore[attr-defined]
+            size_len,
+            core_shape_len,
+            input_bc_patterns,
+            output_bc_patterns,
+            core_cache_key,
+        )
+        random_rv_key = sha256(str(random_rv_key_contents).encode()).hexdigest()
     return random, random_rv_key

{pytensor-2.36.1 → pytensor-2.36.2}/pytensor/link/numba/dispatch/subtensor.py

@@ -649,7 +649,6 @@ def vector_integer_advanced_indexing(
             codegen,
             function_name=func_name,
             global_env=globals(),
-            cache_key=cache_key,
         )
     )
     return ret_func, cache_key

{pytensor-2.36.1 → pytensor-2.36.2}/pytensor/scalar/basic.py

@@ -4300,14 +4300,8 @@ class Composite(ScalarInnerGraphOp):
         self._fgraph = fgraph
         return self._fgraph
 
-    def clone_float32(self):
-        # This will not modify the fgraph or the nodes
-        new_ins, new_outs = composite_f32.apply(self.fgraph)
-        return Composite(new_ins, new_outs)
-
     def clone(self):
-        new_ins, new_outs = composite_f32.apply(self.fgraph)
-        return Composite(new_ins, new_outs)
+        return self.__class__(self.fgraph.inputs, self.fgraph.outputs)
 
     def output_types(self, input_types):
         if tuple(input_types) != self.inputs_type:
@@ -4423,86 +4417,4 @@ class Composite(ScalarInnerGraphOp):
         return self.c_code_template % d
 
     def c_code_cache_version_outer(self) -> tuple[int, ...]:
-        return (6,)
-
-
-class Compositef32:
-    # This is a dict of scalar op classes that need special handling
-    special: dict = {}
-
-    def apply(self, fgraph):
-        mapping = {}
-        topo = fgraph.toposort()
-        for i in fgraph.inputs:
-            if i.dtype == "float16":
-                mapping[i] = get_scalar_type("float32")()
-                if hasattr(i.tag, "test_value"):
-                    mapping[i].tag.test_value = i.tag.test_value
-            else:
-                mapping[i] = i
-        for node in topo:
-            # Patch up for constants
-            for i in node.inputs:
-                if i not in mapping:
-                    assert type(i) is ScalarConstant
-                    if i.type == float16:
-                        ni = ScalarConstant(float32, i.data)
-                    else:
-                        ni = i
-                    mapping[i] = ni
-            if isinstance(node.op, tuple(self.special)):
-                self.special[type(node.op)](node, mapping)
-                continue
-            new_node = node.clone_with_new_inputs(
-                [mapping[inp] for inp in node.inputs], strict=False
-            )
-            # make sure we don't produce any float16.
-            assert not any(o.dtype == "float16" for o in new_node.outputs)
-            mapping.update(zip(node.outputs, new_node.outputs, strict=True))
-
-        new_ins = [mapping[inp] for inp in fgraph.inputs]
-        new_outs = [mapping[out] for out in fgraph.outputs]
-        return new_ins, new_outs
-
-
-composite_f32 = Compositef32()
-
-
-def handle_cast(node, mapping):
-    inp = mapping[node.inputs[0]]
-    out = node.outputs[0]
-    node_ok = False
-    if node.op.o_type == float16:
-        if node.inputs[0].type == float32:
-            # cast f32 -> f16, remove
-            mapping[out] = inp
-            return
-        else:
-            # cast to f16, convert to f32
-            new_out = cast(inp, "float32")
-            # change the node for the following if
-            node = new_out.owner
-            mapping[out] = new_out
-            node_ok = True
-    if node.inputs[0].type == float16:
-        if node.op.o_type == inp.type:
-            # cast f16 to new input type, remove
-            mapping[out] = inp
-            return
-    if not node_ok:
-        new_node = node.clone_with_new_inputs([inp], strict=False)
-        mapping[out] = new_node.outputs[0]
-
-
-Compositef32.special[Cast] = handle_cast
-
-
-def handle_composite(node, mapping):
-    new_op = node.op.clone_float32()
-    new_outs = new_op(*[mapping[i] for i in node.inputs], return_list=True)
-    assert len(new_outs) == len(node.outputs)
-    for o, no in zip(node.outputs, new_outs, strict=True):
-        mapping[o] = no
-
-
-Compositef32.special[Composite] = handle_composite
+        return (7,)

{pytensor-2.36.1 → pytensor-2.36.2}/pytensor/scalar/loop.py

@@ -123,7 +123,7 @@ class ScalarLoop(ScalarInnerGraphOp):
             update, until = self.outputs, None
         init = self.inputs[: len(update)]
         constant = self.inputs[len(update) :]
-        return ScalarLoop(
+        return self.__class__(
             init=init,
             update=update,
             constant=constant,

{pytensor-2.36.1 → pytensor-2.36.2}/pytensor/scan/op.py

@@ -1521,7 +1521,7 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
 
     def clone(self) -> "Scan":
         res = copy(self)
-        res.fgraph = res.fgraph.clone()
+        res.fgraph = res.fgraph.clone(clone_inner_graphs=True)
         return res
 
     def make_thunk(self, node, storage_map, compute_map, no_recycling, impl=None):

{pytensor-2.36.1 → pytensor-2.36.2}/pytensor/tensor/optimize.py

@@ -211,7 +211,7 @@ class ScipyWrapperOp(Op, HasInnerGraph):
 
     def clone(self):
         copy_op = copy(self)
-        copy_op.fgraph = self.fgraph.clone()
+        copy_op.fgraph = self.fgraph.clone(clone_inner_graphs=True)
         return copy_op
 
     def prepare_node(

{pytensor-2.36.1 → pytensor-2.36.2/pytensor.egg-info}/PKG-INFO

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