pytensor 2.36.0__tar.gz → 2.36.2__tar.gz

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

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

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2025-12-15T22:58:11+0100",
+ "date": "2026-01-02T17:20:52+0100",
  "dirty": false,
  "error": null,
- "full-revisionid": "6319fac84ebc71ad079c5986c6c302a52e254198",
- "version": "2.36.0"
+ "full-revisionid": "0f864a33958e68e5110d043c4b5ef549cd1d6074",
+ "version": "2.36.2"
 }
 '''  # END VERSION_JSON
 

{pytensor-2.36.0 → 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.0 → 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.0 → pytensor-2.36.2}/pytensor/link/numba/cache.py

@@ -11,7 +11,7 @@ from pytensor.configdefaults import config
 
 
 NUMBA_CACHE_PATH = config.base_compiledir / "numba"
-NUMBA_CACHE_PATH.mkdir(exist_ok=True)
+NUMBA_CACHE_PATH.mkdir(parents=True, exist_ok=True)
 CACHED_SRC_FUNCTIONS: WeakKeyDictionary[Callable, str] = WeakKeyDictionary()
 
 

{pytensor-2.36.0 → 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.0 → 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.0 → 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.0 → 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.0 → pytensor-2.36.2}/pytensor/link/pytorch/dispatch/scalar.py

@@ -5,6 +5,7 @@ import torch
 from pytensor.link.pytorch.dispatch.basic import pytorch_funcify
 from pytensor.scalar.basic import (
     Cast,
+    Clip,
     Invert,
     ScalarOp,
 )
@@ -71,6 +72,14 @@ def pytorch_funcify_Softplus(op, node, **kwargs):
     return torch.nn.Softplus()
 
 
+@pytorch_funcify.register(Clip)
+def pytorch_funcify_Clip(op, node, **kwargs):
+    def clip(x, min_val, max_val):
+        return torch.where(x < min_val, min_val, torch.where(x > max_val, max_val, x))
+
+    return clip
+
+
 @pytorch_funcify.register(ScalarLoop)
 def pytorch_funicify_ScalarLoop(op, node, **kwargs):
     update = pytorch_funcify(op.fgraph, **kwargs)

{pytensor-2.36.0 → 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.0 → 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.0 → pytensor-2.36.2}/pytensor/scan/basic.py

@@ -1,5 +1,6 @@
 import typing
 import warnings
+from functools import reduce
 from itertools import chain
 
 import numpy as np
@@ -16,7 +17,6 @@ from pytensor.graph.type import HasShape
 from pytensor.graph.utils import MissingInputError, TestValueError
 from pytensor.scan.op import Scan, ScanInfo
 from pytensor.scan.utils import expand_empty, safe_new, until
-from pytensor.tensor.basic import get_underlying_scalar_constant_value
 from pytensor.tensor.exceptions import NotScalarConstantError
 from pytensor.tensor.math import minimum
 from pytensor.tensor.shape import shape_padleft
@@ -143,31 +143,6 @@ def get_updates_and_outputs(ls):
         raise ValueError(error_msg)
 
 
-def isNaN_or_Inf_or_None(x):
-    isNone = x is None
-    try:
-        isNaN = np.isnan(x)
-        isInf = np.isinf(x)
-        isStr = isinstance(x, str)
-    except Exception:
-        isNaN = False
-        isInf = False
-        isStr = False
-    if not isNaN and not isInf:
-        try:
-            val = get_underlying_scalar_constant_value(x)
-            isInf = np.isinf(val)
-            isNaN = np.isnan(val)
-        except Exception:
-            isNaN = False
-            isInf = False
-    if isinstance(x, Constant) and isinstance(x.data, str):
-        isStr = True
-    else:
-        isStr = False
-    return isNone or isNaN or isInf or isStr
-
-
 def _manage_output_api_change(outputs, updates, return_updates):
     if return_updates:
         warnings.warn(
@@ -505,7 +480,7 @@ def scan(
 
     # This helper eagerly skips the Scan if n_steps is known to be 1
     single_step_requested = False
-    if isinstance(n_steps, float | int):
+    if isinstance(n_steps, int | float):
         single_step_requested = n_steps == 1
     else:
         try:
@@ -676,33 +651,20 @@ def scan(
                 if nw_name is not None:
                     nw_seq.name = nw_name
 
-    # Since we've added all sequences now we need to level them up based on
-    # n_steps or their different shapes
-    lengths_vec = [seq.shape[0] for seq in scan_seqs]
-
-    if not isNaN_or_Inf_or_None(n_steps):
-        # ^ N_steps should also be considered
-        lengths_vec.append(pt.as_tensor(n_steps))
-
-    if len(lengths_vec) == 0:
-        # ^ No information about the number of steps
-        raise ValueError(
-            "No information about the number of steps "
-            "provided. Either provide a value for "
-            "n_steps argument of scan or provide an input "
-            "sequence"
-        )
-
-    # If the user has provided the number of steps, do that regardless ( and
-    # raise an error if the sequences are not long enough )
-    if isNaN_or_Inf_or_None(n_steps):
-        actual_n_steps = lengths_vec[0]
-        for contestant in lengths_vec[1:]:
-            actual_n_steps = minimum(actual_n_steps, contestant)
+    if n_steps is None:
+        if not scan_seqs:
+            raise ValueError(
+                "No information about the number of steps provided. "
+                "Either provide a value for n_steps argument of scan or provide an input sequence."
+            )
+        actual_n_steps = reduce(minimum, [seq.shape[0] for seq in scan_seqs])
     else:
-        actual_n_steps = pt.as_tensor(n_steps)
+        actual_n_steps = pt.as_tensor(n_steps, dtype="int64", ndim=0)
 
+    # Since we've added all sequences now we need to level them off based on
+    # n_steps or their different shapes
     scan_seqs = [seq[:actual_n_steps] for seq in scan_seqs]
+
     # Conventions :
     #   mit_mot = multiple input taps, multiple output taps ( only provided
     #             by the gradient function )
@@ -899,10 +861,8 @@ def scan(
         raw_inner_outputs = fn(*args)
 
     condition, outputs, updates = get_updates_and_outputs(raw_inner_outputs)
-    if condition is not None:
-        as_while = True
-    else:
-        as_while = False
+    as_while = condition is not None
+
     ##
     # Step 3. Check if we actually need scan and remove it if we don't
     ##
@@ -934,7 +894,7 @@ def scan(
 
     # extract still missing inputs (there still might be so) and add them
     # as non sequences at the end of our args
-    if condition is not None:
+    if as_while:
         outputs.append(condition)
     fake_nonseqs = [x.type() for x in non_seqs]
     fake_outputs = clone_replace(
@@ -1252,8 +1212,8 @@ def scan(
     rightOrder = (
         mit_sot_rightOrder
         + sit_sot_rightOrder
-        + untraced_sit_sot_rightOrder
         + nit_sot_rightOrder
+        + untraced_sit_sot_rightOrder
     )
     scan_out_list = [None] * len(rightOrder)
     for idx, pos in enumerate(rightOrder):

{pytensor-2.36.0 → 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):