pytensor 2.38.2__tar.gz → 2.38.3__tar.gz

{pytensor-2.38.2/pytensor.egg-info → pytensor-2.38.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pytensor
-Version: 2.38.2
+Version: 2.38.3
 Summary: Optimizing compiler for evaluating mathematical expressions on CPUs and GPUs.
 Author-email: pymc-devs <pymc.devs@gmail.com>
 License-Expression: BSD-3-Clause
@@ -31,7 +31,7 @@ License-File: LICENSE.txt
 Requires-Dist: setuptools>=59.0.0
 Requires-Dist: scipy<2,>=1
 Requires-Dist: numpy>=2.0
-Requires-Dist: numba<1,>0.57
+Requires-Dist: numba<=0.65.1,>=0.58
 Requires-Dist: filelock>=3.15
 Requires-Dist: etuples
 Requires-Dist: logical-unification
@@ -60,7 +60,7 @@ Provides-Extra: jax
 Requires-Dist: jax; extra == "jax"
 Requires-Dist: jaxlib; extra == "jax"
 Provides-Extra: numba
-Requires-Dist: numba>=0.57; extra == "numba"
+Requires-Dist: numba<=0.65.1,>=0.58; extra == "numba"
 Requires-Dist: llvmlite; extra == "numba"
 Dynamic: license-file
 

{pytensor-2.38.2 → pytensor-2.38.3}/pyproject.toml

@@ -50,7 +50,7 @@ dependencies = [
     "setuptools>=59.0.0",
     "scipy>=1,<2",
     "numpy>=2.0",
-    "numba>0.57,<1",
+    "numba>=0.58,<=0.65.1",
     "filelock>=3.15",
     "etuples",
     "logical-unification",
@@ -81,7 +81,7 @@ tests = [
 ]
 rtd = ["sphinx>=5.1.0,<6", "pygments", "pydot"]
 jax = ["jax", "jaxlib"]
-numba = ["numba>=0.57", "llvmlite"]
+numba = ["numba>=0.58,<=0.65.1", "llvmlite"]
 
 [tool.setuptools.packages.find]
 include = ["pytensor*"]
@@ -158,6 +158,7 @@ lines-after-imports = 2
 "pytensor/bin/pytensor_cache.py" = ["T201"]
 "tests/compile/test_monitormode.py" = ["T201"]
 "scripts/run_mypy.py" = ["T201"]
+"scripts/bump_numba_upper_bound.py" = ["T201"]
 # Test modules of optional backends that use `pytest.importorskip`, skip "E402"
 "tests/link/jax/**/test_*.py" = ["E402"]
 "tests/link/numba/**/test_*.py" = ["E402"]

{pytensor-2.38.2 → pytensor-2.38.3}/pytensor/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2026-03-06T13:58:13+0100",
+ "date": "2026-04-27T13:00:34+0200",
  "dirty": false,
  "error": null,
- "full-revisionid": "7fa753b583fd16f46754b356c21055c9559080ef",
- "version": "2.38.2"
+ "full-revisionid": "74cb66529f26343d5034adbf38e8ccd9996449fc",
+ "version": "2.38.3"
 }
 '''  # END VERSION_JSON
 

{pytensor-2.38.2 → pytensor-2.38.3}/pytensor/link/numba/cache.py

@@ -1,3 +1,4 @@
+import uuid
 from collections.abc import Callable
 from hashlib import sha256
 from pickle import dump
@@ -8,6 +9,7 @@ from weakref import WeakKeyDictionary
 import numba
 from llvmlite import ir
 from numba.core import cgutils
+from numba.core.bytecode import FunctionIdentity
 from numba.core.caching import CacheImpl, _CacheLocator
 
 from pytensor.configdefaults import config
@@ -18,6 +20,35 @@ NUMBA_CACHE_PATH.mkdir(parents=True, exist_ok=True)
 CACHED_SRC_FUNCTIONS: WeakKeyDictionary[Callable, str] = WeakKeyDictionary()
 
 
+class _RandomUidIter:
+    """Iterator yielding cryptographically random 128-bit integers.
+
+    Replaces numba's default ``FunctionIdentity._unique_ids`` sequential
+    counter so that sibling forks — which inherit the counter state verbatim —
+    cannot allocate identical uids to same-qualname functions they
+    fresh-compile next. Identical uids produce byte-identical LLVM mangled
+    symbols with different bodies; a later process that loads both hits LLVM's
+    weak-merge and dispatches into the wrong body, segfaulting.
+
+    Numba only consumes ``_unique_ids`` via ``next(cls._unique_ids)`` (see
+    ``numba.core.bytecode.FunctionIdentity.from_function``), so any iterator
+    of unique ints is a drop-in replacement. ``uuid.uuid4`` pulls fresh
+    entropy from the OS on every call, making this generator fork-safe
+    without any re-seed hook.
+
+    See https://github.com/numba/numba/issues/10486.
+    """
+
+    def __iter__(self) -> "_RandomUidIter":
+        return self
+
+    def __next__(self) -> int:
+        return uuid.uuid4().int
+
+
+FunctionIdentity._unique_ids = _RandomUidIter()
+
+
 class NumbaPyTensorCacheLocator(_CacheLocator):
     """Locator for Numba functions defined from PyTensor-generated source code.
 

{pytensor-2.38.2 → pytensor-2.38.3}/pytensor/link/numba/dispatch/basic.py

@@ -13,7 +13,10 @@ from pytensor import config
 from pytensor.graph.basic import Apply, Constant, Variable
 from pytensor.graph.fg import FunctionGraph
 from pytensor.graph.type import Type
-from pytensor.link.numba.cache import compile_numba_function_src, hash_from_pickle_dump
+from pytensor.link.numba.cache import (
+    compile_numba_function_src,
+    hash_from_pickle_dump,
+)
 from pytensor.link.utils import (
     fgraph_to_python,
 )
@@ -33,7 +36,7 @@ def _filter_numba_warnings():
         "ignore",
         message=(
             "(\x1b\\[1m)*"  # ansi escape code for bold text
-            'Cannot cache compiled function "numba_funcified_fgraph" as it uses dynamic globals'
+            'Cannot cache compiled function "numba_funcified_fgraph.*" as it uses dynamic globals'
         ),
         category=NumbaWarning,
     )

{pytensor-2.38.2 → pytensor-2.38.3}/pytensor/tensor/basic.py

@@ -1802,6 +1802,40 @@ def _get_vector_length_Alloc(var_inst, var):
         raise ValueError(f"Length of {var} cannot be determined")
 
 
+@_vectorize_node.register(Alloc)
+def vectorize_alloc(op, node, val, *shape):
+    old_val, *old_shape = node.inputs
+    [old_alloc] = node.outputs
+
+    assert len(shape) == len(old_shape), (
+        "Number of shape entries can't change in vectorize_alloc"
+    )
+
+    if not all(all(s.broadcastable) for s in shape):
+        # May imply a non-square Alloc
+        return vectorize_node_fallback(op, node, val, *shape)
+
+    val_batch_ndim = val.ndim - old_val.ndim
+    shape_batch_ndim = max((s.ndim for s in shape), default=0)
+
+    # Add implicit core dims that alloc prepends (alloc aligns val to the right)
+    n_implicit_core_dims = len(old_shape) - old_val.ndim
+    if n_implicit_core_dims > 0:
+        val = expand_dims(
+            val, list(range(val_batch_ndim, val_batch_ndim + n_implicit_core_dims))
+        )
+
+    new_alloc = alloc(
+        val,
+        *val.shape[:val_batch_ndim],
+        *(s.squeeze() for s in shape),
+    )
+    # Expand leading batch dims implied by the shape entries (if any)
+    new_alloc = atleast_Nd(new_alloc, n=shape_batch_ndim + old_alloc.ndim)
+
+    return [new_alloc]
+
+
 def full(shape, fill_value, dtype=None):
     """Return a new array of given shape and type, filled with `fill_value`.
 

{pytensor-2.38.2 → pytensor-2.38.3}/pytensor/tensor/optimize.py

@@ -35,7 +35,6 @@ from pytensor.tensor.math import tensordot
 from pytensor.tensor.reshape import pack, unpack
 from pytensor.tensor.slinalg import solve
 from pytensor.tensor.variable import TensorVariable, Variable
-from pytensor.utils import unzip
 
 
 # scipy.optimize can be slow to import, and will not be used by most users
@@ -214,22 +213,6 @@ class ScipyWrapperOp(Op, HasInnerGraph):
 
         return Apply(self, inputs, [self.inner_inputs[0].type(), ps.bool("success")])
 
-    def connection_pattern(self, node=None):
-        """
-        All Ops that inherit from ScipyWrapperOp share the same connection pattern logic, because they all share the
-        same output structure. There are two outputs: the optimized variable, and a success flag. The success flag is
-        not differentiable, so it is never connected. The optimized variable is connected only to inputs that are
-        both connected to the objective function and of float dtype.
-        """
-        fgraph = self.fgraph
-        fx = fgraph.outputs[0]
-        return [
-            # Every input is disonnected to the second output (success)
-            # And may or not be connected to the first output (opt_x)
-            [connected, False]
-            for [connected] in io_connection_pattern(fgraph.inputs, [fx])
-        ]
-
 
 class ScipyScalarWrapperOp(ScipyWrapperOp):
     def build_fn(self):
@@ -277,6 +260,7 @@ class ScipyScalarWrapperOp(ScipyWrapperOp):
         inner_fx = self.inner_outputs[0]
 
         if is_minimization:
+            # The implicit function in minimization is grad(x, theta) == 0
             inner_fx = grad(inner_fx, inner_x)
 
         df_dx, *arg_grads = grad(
@@ -287,32 +271,35 @@ class ScipyScalarWrapperOp(ScipyWrapperOp):
             return_disconnected="disconnected",
         )
 
-        outer_arg_grad_map = dict(zip(args, arg_grads))
-        valid_args_and_grads = [
-            (arg, g)
-            for arg, g in outer_arg_grad_map.items()
-            if not isinstance(g.type, DisconnectedType | NullType)
-        ]
+        args_to_diff: tuple[bool, ...] = tuple(
+            not isinstance(g.type, DisconnectedType | NullType) for g in arg_grads
+        )
 
-        if len(valid_args_and_grads) == 0:
+        if not any(args_to_diff):
             # No differentiable arguments, return disconnected gradients
             return arg_grads
 
-        outer_args_to_diff, df_dthetas = unzip(valid_args_and_grads, n=2)
+        df_dthetas = [g for g, to_diff in zip(arg_grads, args_to_diff) if to_diff]
 
-        replace = dict(zip(fgraph.inputs, (x_star, *args), strict=True))
+        # Make gradient an expression of the outer variables
         df_dx_star, *df_dthetas_stars = graph_replace(
-            [df_dx, *df_dthetas], replace=replace
+            [df_dx, *df_dthetas], replace=tuple(zip(fgraph.inputs, (x_star, *args)))
         )
 
-        arg_to_grad = dict(zip(outer_args_to_diff, df_dthetas_stars))
+        grad_wrt_args = []
+        df_dthetas_iter = iter(df_dthetas_stars)
+        for i, (arg, to_diff) in enumerate(zip(args, args_to_diff)):
+            if not to_diff:
+                # Store the null grad we got from the initial `grad` call
+                g = arg_grads[i]
+                assert isinstance(g.type, NullType | DisconnectedType)
+            else:
+                # Compute non-null grad and chain with output_grad
+                df_dtheta_star = next(df_dthetas_iter)
+                g = (-df_dtheta_star / df_dx_star) * output_grad
+            grad_wrt_args.append(g)
 
-        grad_wrt_args = [
-            (-arg_to_grad[arg] / df_dx_star) * output_grad
-            if arg in arg_to_grad
-            else outer_arg_grad_map[arg]
-            for arg in args
-        ]
+        assert next(df_dthetas_iter, None) is None, "Iterator was not exhausted"
 
         return grad_wrt_args
 
@@ -358,7 +345,7 @@ class ScipyVectorWrapperOp(ScipyWrapperOp):
 
         Notes
         -----
-        The gradents are computed using the implicit function theorem. Given a fuction `f(x, theta) = 0`, and a function
+        The gradients are computed using the implicit function theorem. Given a function `f(x, theta) = 0`, and a function
         `x_star(theta)` that, given input parameters theta returns `x` such that `f(x_star(theta), theta) = 0`, we can
         compute the gradients of `x_star` with respect to `theta` as follows:
 
@@ -384,8 +371,12 @@ class ScipyVectorWrapperOp(ScipyWrapperOp):
         fgraph = self.fgraph
         inner_x, *inner_args = self.inner_inputs
         implicit_f = self.inner_outputs[0]
+        if is_minimization:
+            # The implicit function in minimization is grad(x, theta) == 0
+            implicit_f = grad(implicit_f, inner_x)
 
-        df_dx, *arg_grads = grad(
+        # Call grad to see what arguments are connected
+        _, *arg_grads = grad(
             implicit_f.sum(),
             [inner_x, *inner_args],
             disconnected_inputs="ignore",
@@ -393,32 +384,22 @@ class ScipyVectorWrapperOp(ScipyWrapperOp):
             return_disconnected="disconnected",
         )
 
-        inner_args_to_diff = [
-            arg
-            for arg, g in zip(inner_args, arg_grads)
-            if not isinstance(g.type, DisconnectedType | NullType)
-        ]
+        args_to_diff: tuple[bool, ...] = tuple(
+            not isinstance(g.type, DisconnectedType | NullType) for g in arg_grads
+        )
 
-        if len(inner_args_to_diff) == 0:
+        if not args_to_diff:
             # No differentiable arguments, return disconnected/null gradients
             return arg_grads
 
-        outer_args_to_diff = [
-            arg
-            for inner_arg, arg in zip(inner_args, args)
-            if inner_arg in inner_args_to_diff
-        ]
-        invalid_grad_map = {
-            arg: g for arg, g in zip(args, arg_grads) if arg not in outer_args_to_diff
-        }
-
-        if is_minimization:
-            implicit_f = grad(implicit_f, inner_x)
-
         # Gradients are computed using the inner graph of the optimization op, not the actual inputs/outputs of the op.
         packed_inner_args, packed_arg_shapes, implicit_f = pack_inputs_of_objective(
             implicit_f,
-            inner_args_to_diff,
+            [
+                inner_arg
+                for inner_arg, to_diff in zip(inner_args, args_to_diff)
+                if to_diff
+            ],
         )
 
         df_dx, df_dtheta = jacobian(
@@ -430,11 +411,11 @@ class ScipyVectorWrapperOp(ScipyWrapperOp):
 
         # Replace inner inputs (abstract dummies) with outer inputs (the actual user-provided symbols)
         # at the solution point. Innner arguments aren't needed anymore, delete them to avoid accidental references.
-        del inner_x
-        del inner_args
-        inner_to_outer_map = dict(zip(fgraph.inputs, (x_star, *args)))
+        del inner_x, inner_args
+        inner_to_outer_map = tuple(zip(fgraph.inputs, (x_star, *args)))
         df_dx_star, df_dtheta_star = graph_replace(
-            [df_dx, df_dtheta], inner_to_outer_map
+            [df_dx, df_dtheta],
+            replace=inner_to_outer_map,
         )
 
         if df_dtheta_star.ndim == 0 or df_dx_star.ndim == 0:
@@ -454,27 +435,49 @@ class ScipyVectorWrapperOp(ScipyWrapperOp):
         else:
             grad_wrt_args = [grad_wrt_args_packed]
 
-        arg_to_grad = dict(zip(outer_args_to_diff, grad_wrt_args))
-
         final_grads = []
-        for arg in args:
-            arg_grad = arg_to_grad.get(arg, None)
-
-            if arg_grad is None:
-                final_grads.append(invalid_grad_map[arg])
-                continue
-
-            if arg_grad.ndim > 0 and output_grad.ndim > 0:
-                g = tensordot(output_grad, arg_grad, [[0], [0]])
+        grad_wrt_args_iter = iter(grad_wrt_args)
+        for i, (arg, to_diff) in enumerate(zip(args, args_to_diff)):
+            if not to_diff:
+                # Store the null grad we got from the initial `grad` call
+                g = arg_grads[i]
+                assert isinstance(g.type, NullType | DisconnectedType)
             else:
-                g = arg_grad * output_grad
-            if isinstance(arg.type, ScalarType) and isinstance(g, TensorVariable):
-                g = scalar_from_tensor(g)
+                # Compute non-null grad and chain with output_grad
+                arg_grad = next(grad_wrt_args_iter)
+                if arg_grad.ndim > 0 and output_grad.ndim > 0:
+                    g = tensordot(output_grad, arg_grad, [[0], [0]])
+                else:
+                    g = arg_grad * output_grad
+                if isinstance(arg.type, ScalarType) and isinstance(g, TensorVariable):
+                    g = scalar_from_tensor(g)
             final_grads.append(g)
 
+        assert next(grad_wrt_args_iter, None) is None, "Iterator was not exhausted"
+
         return final_grads
 
 
+def _optimizer_connection_pattern(fgraph, is_minimization):
+    """Compute connection pattern for scipy optimization Ops.
+
+    There are two outputs: the optimized variable, and a success flag. The success flag is
+    not differentiable, so it is never connected. The optimized variable is connected only
+    to inputs that are connected to the implicit function used in the gradient computation.
+
+    For minimization, the implicit function is grad(objective, x), not the objective itself.
+    An input may be connected to the objective but disconnected from its gradient (e.g. an
+    additive constant), so the connection pattern must reflect the actual implicit function.
+    """
+    inner_x = fgraph.inputs[0]
+    fx = fgraph.outputs[0]
+    if is_minimization:
+        fx = grad(fx, inner_x)
+    return [
+        [connected, False] for [connected] in io_connection_pattern(fgraph.inputs, [fx])
+    ]
+
+
 class MinimizeScalarOp(ScipyScalarWrapperOp):
     def __init__(
         self,
@@ -503,6 +506,9 @@ class MinimizeScalarOp(ScipyScalarWrapperOp):
         self._fn = None
         self._fn_wrapped = None
 
+    def connection_pattern(self, node=None):
+        return _optimizer_connection_pattern(self.fgraph, is_minimization=True)
+
     def __str__(self):
         return f"{self.__class__.__name__}(method={self.method})"
 
@@ -639,6 +645,9 @@ class MinimizeOp(ScipyVectorWrapperOp):
         self._fn = None
         self._fn_wrapped = None
 
+    def connection_pattern(self, node=None):
+        return _optimizer_connection_pattern(self.fgraph, is_minimization=True)
+
     def __str__(self):
         str_args = ", ".join(
             [
@@ -834,6 +843,9 @@ class RootScalarOp(ScipyScalarWrapperOp):
         self._fn = None
         self._fn_wrapped = None
 
+    def connection_pattern(self, node=None):
+        return _optimizer_connection_pattern(self.fgraph, is_minimization=False)
+
     def __str__(self):
         str_args = ", ".join(
             [f"{arg}={getattr(self, arg)}" for arg in ["method", "jac", "hess"]]
@@ -980,6 +992,9 @@ class RootOp(ScipyVectorWrapperOp):
         self._fn = None
         self._fn_wrapped = None
 
+    def connection_pattern(self, node=None):
+        return _optimizer_connection_pattern(self.fgraph, is_minimization=False)
+
     def __str__(self):
         str_args = ", ".join(
             [f"{arg}={getattr(self, arg)}" for arg in ["method", "jac"]]

{pytensor-2.38.2 → pytensor-2.38.3}/pytensor/tensor/rewriting/blockwise.py

@@ -5,7 +5,14 @@ from pytensor.graph.replace import vectorize_graph
 from pytensor.graph.rewriting.basic import copy_stack_trace, dfs_rewriter
 from pytensor.graph.rewriting.unify import OpPattern, OpPatternOpTypeType
 from pytensor.graph.traversal import apply_ancestors
-from pytensor.tensor.basic import Alloc, ARange, alloc, shape_padleft
+from pytensor.tensor.basic import (
+    Alloc,
+    AllocEmpty,
+    ARange,
+    alloc,
+    expand_dims,
+    shape_padleft,
+)
 from pytensor.tensor.blockwise import Blockwise, _squeeze_left
 from pytensor.tensor.math import Dot
 from pytensor.tensor.rewriting.basic import (
@@ -89,6 +96,7 @@ optdb.register(
         blockwise_of(
             Dot
             | Alloc
+            | AllocEmpty
             | ARange
             | Subtensor
             | AdvancedSubtensor
@@ -106,7 +114,7 @@ def local_eager_useless_unbatched_blockwise(fgraph, node):
 
 @register_specialize("shape_unsafe")
 @node_rewriter([Blockwise])
-def local_blockwise_alloc(fgraph, node):
+def local_blockwise_alloc_inputs(fgraph, node):
     """Push Allocs from the inputs to the output of Blockwise Ops.
 
     BOp = Blockwise(Op, signature="(x),(x)->(x)")
@@ -218,6 +226,34 @@ def local_blockwise_alloc(fgraph, node):
     return new_outs
 
 
+@register_canonicalize
+@register_specialize
+@node_rewriter([blockwise_of(Alloc)])
+def local_blockwise_alloc(fgraph, node):
+    val, *shape = node.inputs
+
+    if not all(all(s.broadcastable) for s in shape):
+        # May imply a non-square Alloc
+        return None
+
+    batch_ndim = node.op.batch_ndim(node)
+
+    # Add implicit core dims that alloc prepends (alloc aligns val to the right)
+    n_implicit_core_dims = node.outputs[0].ndim - val.ndim
+    if n_implicit_core_dims > 0:
+        val = expand_dims(
+            val, list(range(batch_ndim, batch_ndim + n_implicit_core_dims))
+        )
+
+    new_alloc = alloc(
+        val,
+        *val.shape[:batch_ndim],
+        *(s.squeeze() for s in shape),
+    )
+    copy_stack_trace(node.outputs[0], new_alloc)
+    return [new_alloc]
+
+
 @register_specialize
 @node_rewriter([blockwise_of(Reshape)])
 def local_blockwise_reshape(fgraph, node):

{pytensor-2.38.2 → pytensor-2.38.3/pytensor.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pytensor
-Version: 2.38.2
+Version: 2.38.3
 Summary: Optimizing compiler for evaluating mathematical expressions on CPUs and GPUs.
 Author-email: pymc-devs <pymc.devs@gmail.com>
 License-Expression: BSD-3-Clause
@@ -31,7 +31,7 @@ License-File: LICENSE.txt
 Requires-Dist: setuptools>=59.0.0
 Requires-Dist: scipy<2,>=1
 Requires-Dist: numpy>=2.0
-Requires-Dist: numba<1,>0.57
+Requires-Dist: numba<=0.65.1,>=0.58
 Requires-Dist: filelock>=3.15
 Requires-Dist: etuples
 Requires-Dist: logical-unification
@@ -60,7 +60,7 @@ Provides-Extra: jax
 Requires-Dist: jax; extra == "jax"
 Requires-Dist: jaxlib; extra == "jax"
 Provides-Extra: numba
-Requires-Dist: numba>=0.57; extra == "numba"
+Requires-Dist: numba<=0.65.1,>=0.58; extra == "numba"
 Requires-Dist: llvmlite; extra == "numba"
 Dynamic: license-file
 

{pytensor-2.38.2 → pytensor-2.38.3}/pytensor.egg-info/requires.txt

@@ -1,7 +1,7 @@
 setuptools>=59.0.0
 scipy<2,>=1
 numpy>=2.0
-numba<1,>0.57
+numba<=0.65.1,>=0.58
 filelock>=3.15
 etuples
 logical-unification
@@ -22,7 +22,7 @@ jax
 jaxlib
 
 [numba]
-numba>=0.57
+numba<=0.65.1,>=0.58
 llvmlite
 
 [rtd]