pytensor 2.36.2__tar.gz → 2.37.0__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pytensor
-Version: 2.36.2
+Version: 2.37.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.36.2 → pytensor-2.37.0}/doc/extending/creating_an_op.rst

@@ -495,7 +495,7 @@ If both outputs are disconnected PyTensor will not bother calling the :meth:`L_o
 
     from pytensor.graph.op import Op
     from pytensor.graph.basic import Apply
-    from pytensor.gradient import DisconnectedType
+    from pytensor.gradient import DisconnectedType, disconnected_type
 
     class TransposeAndSumOp(Op):
         __props__ = ()
@@ -539,13 +539,13 @@ If both outputs are disconnected PyTensor will not bother calling the :meth:`L_o
             out1_grad, out2_grad = output_grads
 
             if isinstance(out1_grad.type, DisconnectedType):
-                x_grad = DisconnectedType()()
+                x_grad = disconnected_type()
             else:
                 # Transpose the last two dimensions of the output gradient
                 x_grad = pt.swapaxes(out1_grad, -1, -2)
 
             if isinstance(out2_grad.type, DisconnectedType):
-                y_grad = DisconnectedType()()
+                y_grad = disconnected_type()
             else:
                 # Broadcast the output gradient to the same shape as y
                 y_grad = pt.broadcast_to(pt.expand_dims(out2_grad, -1), y.shape)

{pytensor-2.36.2 → pytensor-2.37.0}/pyproject.toml

@@ -149,27 +149,21 @@ lines-after-imports = 2
 # TODO: Get rid of these:
 "**/__init__.py" = ["F401", "E402", "F403"]
 "pytensor/tensor/linalg.py" = ["F403"]
+# Modules that use print-statements, skip "T201"
 "pytensor/link/c/cmodule.py" = ["PTH", "T201"]
 "pytensor/misc/elemwise_time_test.py" = ["T201"]
 "pytensor/misc/elemwise_openmp_speedup.py" = ["T201"]
 "pytensor/misc/check_duplicate_key.py" = ["T201"]
 "pytensor/misc/check_blas.py" = ["T201"]
 "pytensor/bin/pytensor_cache.py" = ["T201"]
-# For the tests we skip `E402` because `pytest.importorskip` is used:
-"tests/link/jax/test_scalar.py" = ["E402"]
-"tests/link/jax/test_tensor_basic.py" = ["E402"]
-"tests/link/numba/test_basic.py" = ["E402"]
-"tests/link/numba/test_cython_support.py" = ["E402"]
-"tests/link/numba/test_performance.py" = ["E402"]
-"tests/link/numba/test_sparse.py" = ["E402"]
-"tests/link/numba/test_tensor_basic.py" = ["E402"]
-"tests/tensor/test_math_scipy.py" = ["E402"]
-"tests/sparse/test_basic.py" = ["E402"]
-"tests/sparse/test_sp2.py" = ["E402"]
-"tests/sparse/test_utils.py" = ["E402"]
-"tests/sparse/sandbox/test_sp.py" = ["E402", "F401"]
 "tests/compile/test_monitormode.py" = ["T201"]
 "scripts/run_mypy.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"]
+"tests/link/pytorch/**/test_*.py" = ["E402"]
+"tests/link/mlx/**/test_*.py" = ["E402"]
+
 
 
 [tool.mypy]

{pytensor-2.36.2 → pytensor-2.37.0}/pytensor/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2026-01-02T17:20:52+0100",
+ "date": "2026-01-23T09:04:26+0100",
  "dirty": false,
  "error": null,
- "full-revisionid": "0f864a33958e68e5110d043c4b5ef549cd1d6074",
- "version": "2.36.2"
+ "full-revisionid": "9834e96cee86d9f18436be5fa9e31c692bf6326f",
+ "version": "2.37.0"
 }
 '''  # END VERSION_JSON
 

{pytensor-2.36.2 → pytensor-2.37.0}/pytensor/breakpoint.py

@@ -1,6 +1,6 @@
 import numpy as np
 
-from pytensor.gradient import DisconnectedType
+from pytensor.gradient import disconnected_type
 from pytensor.graph.basic import Apply, Variable
 from pytensor.graph.op import Op
 from pytensor.tensor.basic import as_tensor_variable
@@ -142,7 +142,7 @@ class PdbBreakpoint(Op):
                 output_storage[i][0] = inputs[i + 1]
 
     def grad(self, inputs, output_gradients):
-        return [DisconnectedType()(), *output_gradients]
+        return [disconnected_type(), *output_gradients]
 
     def infer_shape(self, fgraph, inputs, input_shapes):
         # Return the shape of every input but the condition (first input)

{pytensor-2.36.2 → pytensor-2.37.0}/pytensor/configdefaults.py

@@ -270,15 +270,6 @@ def add_basic_configvars():
         in_c_key=False,
     )
 
-    config.add(
-        "conv__assert_shape",
-        "If True, AbstractConv* ops will verify that user-provided"
-        " shapes match the runtime shapes (debugging option,"
-        " may slow down compilation)",
-        BoolParam(False),
-        in_c_key=False,
-    )
-
     config.add(
         "print_global_stats",
         "Print some global statistics (time spent) at the end",

{pytensor-2.36.2 → pytensor-2.37.0}/pytensor/configparser.py

@@ -71,7 +71,6 @@ class PyTensorConfigParser:
     pickle_test_value: bool
     cast_policy: str
     device: str
-    conv__assert_shape: bool
     print_global_stats: bool
     unpickle_function: bool
     # add_compile_configvars

{pytensor-2.36.2 → pytensor-2.37.0}/pytensor/gradient.py

@@ -494,22 +494,25 @@ def Lop(
         coordinates of the tensor elements.
         If `f` is a list/tuple, then return a list/tuple with the results.
     """
-    if not isinstance(eval_points, list | tuple):
-        _eval_points: list[Variable] = [pytensor.tensor.as_tensor_variable(eval_points)]
-    else:
-        _eval_points = [pytensor.tensor.as_tensor_variable(x) for x in eval_points]
+    from pytensor.tensor import as_tensor_variable
 
-    if not isinstance(f, list | tuple):
-        _f: list[Variable] = [pytensor.tensor.as_tensor_variable(f)]
-    else:
-        _f = [pytensor.tensor.as_tensor_variable(x) for x in f]
+    if not isinstance(eval_points, Sequence):
+        eval_points = [eval_points]
+    _eval_points = [
+        x if isinstance(x, Variable) else as_tensor_variable(x) for x in eval_points
+    ]
+
+    if not isinstance(f, Sequence):
+        f = [f]
+    _f = [x if isinstance(x, Variable) else as_tensor_variable(x) for x in f]
 
     grads = list(_eval_points)
 
-    if not isinstance(wrt, list | tuple):
-        _wrt: list[Variable] = [pytensor.tensor.as_tensor_variable(wrt)]
-    else:
-        _wrt = [pytensor.tensor.as_tensor_variable(x) for x in wrt]
+    using_list = isinstance(wrt, list)
+    using_tuple = isinstance(wrt, tuple)
+    if not isinstance(wrt, Sequence):
+        wrt = [wrt]
+    _wrt = [x if isinstance(x, Variable) else as_tensor_variable(x) for x in wrt]
 
     assert len(_f) == len(grads)
     known = dict(zip(_f, grads, strict=True))
@@ -523,8 +526,6 @@ def Lop(
         return_disconnected=return_disconnected,
     )
 
-    using_list = isinstance(wrt, list)
-    using_tuple = isinstance(wrt, tuple)
     return as_list_or_tuple(using_list, using_tuple, ret)
 
 
@@ -2104,16 +2105,13 @@ def jacobian(
             idx, expr, *wrt = args
             return grad(expr[idx], wrt, **grad_kwargs)
 
-        jacobian_matrices, updates = pytensor.scan(
+        jacobian_matrices = pytensor.scan(
             inner_function,
             sequences=pytensor.tensor.arange(expression.size),
             non_sequences=[expression.ravel(), *wrt],
+            return_updates=False,
             return_list=True,
         )
-        if updates:
-            raise ValueError(
-                "The scan used to build the jacobian matrices returned a list of updates"
-            )
 
     if jacobian_matrices[0].ndim < (expression.ndim + wrt[0].ndim):
         # There was some raveling or squeezing done prior to getting the jacobians

{pytensor-2.36.2 → pytensor-2.37.0}/pytensor/link/basic.py

@@ -284,7 +284,7 @@ class PerformLinker(LocalLinker):
     """
 
     required_rewrites: tuple[str, ...] = ("minimum_compile", "py_only")
-    incompatible_rewrites: tuple[str, ...] = ("cxx",)
+    incompatible_rewrites: tuple[str, ...] = ("cxx_only",)
 
     def __init__(
         self, allow_gc: bool | None = None, schedule: Callable | None = None

{pytensor-2.36.2 → pytensor-2.37.0}/pytensor/link/jax/dispatch/scan.py

@@ -10,7 +10,7 @@ from pytensor.scan.op import Scan
 
 
 @jax_funcify.register(Scan)
-def jax_funcify_Scan(op: Scan, **kwargs):
+def jax_funcify_Scan(op: Scan, node, **kwargs):
     # Note: This implementation is different from the internal PyTensor Scan op.
     # In particular, we don't make use of the provided buffers for recurring outputs (MIT-SOT, SIT-SOT)
     # These buffers include the initial state and enough space to store as many intermediate results as needed.
@@ -29,7 +29,10 @@ def jax_funcify_Scan(op: Scan, **kwargs):
 
     # Optimize inner graph (exclude any defalut rewrites that are incompatible with JAX mode)
     rewriter = (
-        get_mode(op.mode).including("jax").excluding(*JAX._optimizer.exclude).optimizer
+        get_mode(op.mode)
+        .including("jax")
+        .excluding("numba", *JAX._optimizer.exclude)
+        .optimizer
     )
     rewriter(op.fgraph)
     scan_inner_func = jax_funcify(op.fgraph, **kwargs)
@@ -216,6 +219,8 @@ def jax_funcify_Scan(op: Scan, **kwargs):
                     if trace.shape[0] > buffer_size:
                         # Trace is longer than buffer, keep just the last `buffer.shape[0]` entries
                         partial_trace = trace[-buffer_size:]
+                    elif trace.shape[0] == buffer_size:
+                        partial_trace = trace
                     else:
                         # Trace is shorter than buffer, this happens when we keep the initial_state
                         if init_state.ndim < buffer.ndim:

{pytensor-2.36.2 → pytensor-2.37.0}/pytensor/link/jax/dispatch/slinalg.py

@@ -3,6 +3,7 @@ import warnings
 import jax
 
 from pytensor.link.jax.dispatch.basic import jax_funcify
+from pytensor.tensor._linalg.solve.linear_control import SolveSylvester
 from pytensor.tensor.slinalg import (
     LU,
     QR,
@@ -13,6 +14,7 @@ from pytensor.tensor.slinalg import (
     Expm,
     LUFactor,
     PivotToPermutations,
+    Schur,
     Solve,
     SolveTriangular,
 )
@@ -92,7 +94,6 @@ def jax_funcify_Solve(op, **kwargs):
 def jax_funcify_SolveTriangular(op, **kwargs):
     lower = op.lower
     unit_diagonal = op.unit_diagonal
-    check_finite = op.check_finite
 
     def solve_triangular(A, b):
         return jax.scipy.linalg.solve_triangular(
@@ -101,7 +102,7 @@ def jax_funcify_SolveTriangular(op, **kwargs):
             lower=lower,
             trans=0,  # this is handled by explicitly transposing A, so it will always be 0 when we get to here.
             unit_diagonal=unit_diagonal,
-            check_finite=check_finite,
+            check_finite=False,
         )
 
     return solve_triangular
@@ -132,27 +133,23 @@ def jax_funcify_PivotToPermutation(op, **kwargs):
 def jax_funcify_LU(op, **kwargs):
     permute_l = op.permute_l
     p_indices = op.p_indices
-    check_finite = op.check_finite
 
     if p_indices:
         raise ValueError("JAX does not support the p_indices argument")
 
     def lu(*inputs):
-        return jax.scipy.linalg.lu(
-            *inputs, permute_l=permute_l, check_finite=check_finite
-        )
+        return jax.scipy.linalg.lu(*inputs, permute_l=permute_l, check_finite=False)
 
     return lu
 
 
 @jax_funcify.register(LUFactor)
 def jax_funcify_LUFactor(op, **kwargs):
-    check_finite = op.check_finite
     overwrite_a = op.overwrite_a
 
     def lu_factor(a):
         return jax.scipy.linalg.lu_factor(
-            a, check_finite=check_finite, overwrite_a=overwrite_a
+            a, check_finite=False, overwrite_a=overwrite_a
         )
 
     return lu_factor
@@ -161,12 +158,11 @@ def jax_funcify_LUFactor(op, **kwargs):
 @jax_funcify.register(CholeskySolve)
 def jax_funcify_ChoSolve(op, **kwargs):
     lower = op.lower
-    check_finite = op.check_finite
     overwrite_b = op.overwrite_b
 
     def cho_solve(c, b):
         return jax.scipy.linalg.cho_solve(
-            (c, lower), b, check_finite=check_finite, overwrite_b=overwrite_b
+            (c, lower), b, check_finite=False, overwrite_b=overwrite_b
         )
 
     return cho_solve
@@ -189,3 +185,27 @@ def jax_funcify_Expm(op, **kwargs):
         return jax.scipy.linalg.expm(x)
 
     return expm
+
+
+@jax_funcify.register(Schur)
+def jax_funcify_Schur(op, **kwargs):
+    output = op.output
+
+    if op.sort is not None:
+        warnings.warn(
+            "jax.scipy.linalg.schur only supports sort=None. The sort argument is ignored."
+        )
+
+    def schur(a):
+        T, Z = jax.scipy.linalg.schur(a, output=output)
+        return T, Z
+
+    return schur
+
+
+@jax_funcify.register(SolveSylvester)
+def jax_funcify_SolveSylsterer(op, **kwargs):
+    def solve_sylvester(a, b, c):
+        return jax.scipy.linalg.solve_sylvester(a, b, c)
+
+    return solve_sylvester

{pytensor-2.36.2 → pytensor-2.37.0}/pytensor/link/jax/dispatch/tensor_basic.py

@@ -17,16 +17,15 @@ from pytensor.tensor.basic import (
     ScalarFromTensor,
     Split,
     TensorFromScalar,
-    Tri,
     get_scalar_constant_value,
 )
 from pytensor.tensor.exceptions import NotScalarConstantError
 from pytensor.tensor.shape import Shape_i
 
 
-ARANGE_CONCRETE_VALUE_ERROR = """JAX requires the arguments of `jax.numpy.arange`
-to be constants. The graph that you defined thus cannot be JIT-compiled
-by JAX. An example of a graph that can be compiled to JAX:
+ARANGE_CONCRETE_VALUE_ERROR = """JAX requires the arguments of `jax.numpy.arange` to be constants.
+The graph that you defined thus cannot be JIT-compiled by JAX.
+An example of a graph that can be compiled to JAX:
 >>> import pytensor.tensor as pt
 >>> pt.arange(1, 10, 2)
 """
@@ -185,19 +184,3 @@ def jax_funcify_ScalarFromTensor(op, **kwargs):
         return jnp.array(x).flatten()[0]
 
     return scalar_from_tensor
-
-
-@jax_funcify.register(Tri)
-def jax_funcify_Tri(op, node, **kwargs):
-    # node.inputs is N, M, k
-    const_args = [getattr(x, "data", None) for x in node.inputs]
-
-    def tri(*args):
-        # args is N, M, k
-        args = [
-            x if const_x is None else const_x
-            for x, const_x in zip(args, const_args, strict=True)
-        ]
-        return jnp.tri(*args, dtype=op.dtype)
-
-    return tri

{pytensor-2.36.2 → pytensor-2.37.0}/pytensor/link/jax/linker.py

@@ -14,7 +14,7 @@ class JAXLinker(JITLinker):
         "jax",
     )  # TODO: Distinguish between optional "jax" and "minimum_compile_jax"
     incompatible_rewrites = (
-        "cxx",
+        "cxx_only",
         "BlasOpt",
         "local_careduce_fusion",
         "scan_save_mem_prealloc",

{pytensor-2.36.2 → pytensor-2.37.0}/pytensor/link/mlx/dispatch/core.py

@@ -13,7 +13,6 @@ from pytensor.tensor.basic import (
     ScalarFromTensor,
     Split,
     TensorFromScalar,
-    Tri,
     get_scalar_constant_value,
 )
 from pytensor.tensor.exceptions import NotScalarConstantError
@@ -219,23 +218,6 @@ def mlx_funcify_ScalarFromTensor(op, **kwargs):
     return scalar_from_tensor
 
 
-@mlx_funcify.register(Tri)
-def mlx_funcify_Tri(op, node, **kwargs):
-    # node.inputs  ->  N, M, k
-    const_args = [getattr(inp, "data", None) for inp in node.inputs]
-    dtype = convert_dtype_to_mlx(op.dtype)
-
-    def tri(*args):
-        # Replace args with compile-time constants when available
-        args = [
-            arg if const_a is None else const_a
-            for arg, const_a in zip(args, const_args, strict=True)
-        ]
-        return mx.tri(*args, dtype=dtype)
-
-    return tri
-
-
 @mlx_funcify.register(AllocEmpty)
 def mlx_funcify_AllocEmpty(op, node, **kwargs):
     dtype = convert_dtype_to_mlx(op.dtype)

{pytensor-2.36.2 → pytensor-2.37.0}/pytensor/link/numba/cache.py

@@ -5,6 +5,9 @@ from tempfile import NamedTemporaryFile
 from typing import Any
 from weakref import WeakKeyDictionary
 
+import numba
+from llvmlite import ir
+from numba.core import cgutils
 from numba.core.caching import CacheImpl, _CacheLocator
 
 from pytensor.configdefaults import config
@@ -127,3 +130,49 @@ def compile_numba_function_src(
         CACHED_SRC_FUNCTIONS[res] = cache_key
 
     return res  # type: ignore
+
+
+@numba.extending.intrinsic(prefer_literal=True)
+def _call_cached_ptr(typingctx, get_ptr_func, func_type_ref, unique_func_name_lit):
+    """
+    Enable caching of function pointers returned by `get_ptr_func`.
+
+    When one of our Numba-dispatched functions depends on a pointer to a compiled function function (e.g. when we call
+    cython_lapack routines), numba will refuse to cache the function, because the pointer may change between runs.
+
+    This intrinsic allows us to cache the pointer ourselves, by storing it in a global variable keyed by a literal
+    `unique_func_name_lit`. The first time the intrinsic is called, it will call `get_ptr_func` to get the pointer, store it
+    in the global variable, and return it. Subsequent calls will load the pointer from the global variable.
+    """
+    func_type = func_type_ref.instance_type
+    cache_key = unique_func_name_lit.literal_value
+
+    def codegen(context, builder, signature, args):
+        ptr_ty = ir.PointerType(ir.IntType(8))
+        null = ptr_ty(None)
+        align = 64
+
+        mod = builder.module
+        var = cgutils.add_global_variable(mod, ptr_ty, f"_ptr_cache_{cache_key}")
+        var.align = align
+        var.linkage = "private"
+        var.initializer = null
+
+        var_val = builder.load_atomic(var, "acquire", align)
+        result_ptr = cgutils.alloca_once_value(builder, var_val)
+
+        with builder.if_then(builder.icmp_signed("==", var_val, null), likely=False):
+            sig = typingctx.resolve_function_type(get_ptr_func, [], {})
+            f = context.get_function(get_ptr_func, sig)
+            new_ptr = f(builder, [])
+            new_ptr = builder.inttoptr(new_ptr, ptr_ty)
+            builder.store_atomic(new_ptr, var, "release", align)
+            builder.store(new_ptr, result_ptr)
+
+        sfunc = cgutils.create_struct_proxy(func_type)(context, builder)
+        sfunc.c_addr = builder.load(result_ptr)
+
+        return sfunc._getvalue()
+
+    sig = func_type(get_ptr_func, func_type_ref, unique_func_name_lit)
+    return sig, codegen

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

@@ -14,14 +14,13 @@ 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.dispatch.sparse import CSCMatrixType, CSRMatrixType
 from pytensor.link.utils import (
     fgraph_to_python,
 )
 from pytensor.scalar.basic import ScalarType
 from pytensor.sparse import SparseTensorType
 from pytensor.tensor.random.type import RandomGeneratorType
-from pytensor.tensor.type import TensorType
+from pytensor.tensor.type import DenseTensorType
 from pytensor.tensor.utils import hash_from_ndarray
 from pytensor.typed_list import TypedListType
 
@@ -112,7 +111,7 @@ def get_numba_type(
         Return Numba scalars for zero dimensional :class:`TensorType`\s.
     """
 
-    if isinstance(pytensor_type, TensorType):
+    if isinstance(pytensor_type, DenseTensorType):
         dtype = pytensor_type.numpy_dtype
         numba_dtype = numba.from_dtype(dtype)
         if force_scalar or (
@@ -125,18 +124,26 @@ def get_numba_type(
         numba_dtype = numba.from_dtype(dtype)
         return numba_dtype
     elif isinstance(pytensor_type, SparseTensorType):
-        dtype = pytensor_type.numpy_dtype
-        numba_dtype = numba.from_dtype(dtype)
+        from pytensor.link.numba.dispatch.sparse.variable import (
+            CSCMatrixType,
+            CSRMatrixType,
+        )
+
+        data_array = numba.types.Array(
+            numba.from_dtype(pytensor_type.numpy_dtype), 1, layout
+        )
+        indices_array = numba.types.Array(numba.from_dtype(np.int32), 1, layout)
+        indptr_array = numba.types.Array(numba.from_dtype(np.int32), 1, layout)
         if pytensor_type.format == "csr":
-            return CSRMatrixType(numba_dtype)
+            return CSRMatrixType(data_array, indices_array, indptr_array)
         if pytensor_type.format == "csc":
-            return CSCMatrixType(numba_dtype)
+            return CSCMatrixType(data_array, indices_array, indptr_array)
     elif isinstance(pytensor_type, RandomGeneratorType):
         return numba.types.NumPyRandomGeneratorType("NumPyRandomGeneratorType")
     elif isinstance(pytensor_type, TypedListType):
         return numba.types.List(get_numba_type(pytensor_type.ttype))
-    else:
-        raise NotImplementedError(f"Numba type not implemented for {pytensor_type}")
+
+    raise NotImplementedError(f"Numba type not implemented for {pytensor_type}")
 
 
 def create_numba_signature(
@@ -171,40 +178,6 @@ def create_numba_signature(
         return numba.types.void(*input_types)
 
 
-def create_tuple_creator(f, n):
-    """Construct a compile-time ``tuple``-comprehension-like loop.
-
-    See https://github.com/numba/numba/issues/2771#issuecomment-414358902
-    """
-    warnings.warn(
-        "create_tuple_creator is deprecated and will be removed in a future release",
-        FutureWarning,
-    )
-
-    assert n > 0
-
-    f = numba_njit(f)
-
-    @numba_njit
-    def creator(args):
-        return (f(0, *args),)
-
-    for i in range(1, n):
-
-        @numba_njit
-        def creator(args, creator=creator, i=i):
-            return (*creator(args), f(i, *args))
-
-    return numba_njit(lambda *args: creator(args))
-
-
-def create_tuple_string(x):
-    if len(x) == 1:
-        return f"({x[0]},)"
-    else:
-        return f"({', '.join(x)})"
-
-
 @numba.extending.intrinsic
 def direct_cast(typingctx, val, typ):
     if isinstance(typ, numba.types.TypeRef):

{pytensor-2.36.2 → pytensor-2.37.0}/pytensor/link/numba/dispatch/blockwise.py

@@ -86,6 +86,7 @@ def numba_funcify_Blockwise(op: BlockwiseWithCoreShape, node, **kwargs):
                 output_bc_patterns,
                 output_dtypes,
                 inplace_pattern,
+                False,  # allow_core_scalar
                 (),  # constant_inputs
                 inputs,
                 tuple_core_shapes,
@@ -98,6 +99,7 @@ def numba_funcify_Blockwise(op: BlockwiseWithCoreShape, node, **kwargs):
         # If the core op cannot be cached, the Blockwise wrapper cannot be cached either
         blockwise_key = None
     else:
+        blockwise_cache_version = 1
         blockwise_key = "_".join(
             map(
                 str,
@@ -108,6 +110,7 @@ def numba_funcify_Blockwise(op: BlockwiseWithCoreShape, node, **kwargs):
                     blockwise_op.signature,
                     input_bc_patterns,
                     core_op_key,
+                    blockwise_cache_version,
                 ),
             )
         )

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

@@ -68,7 +68,7 @@ def numba_funcify_OpFromGraph(op, node=None, **kwargs):
     output_specs = [Out(o, borrow=False) for o in fgraph.outputs]
     insert_deepcopy(fgraph, wrapped_inputs=input_specs, wrapped_outputs=output_specs)
     fgraph_fn, fgraph_cache_key = numba_funcify_and_cache_key(
-        fgraph, squeeze_output=True, **kwargs
+        fgraph, squeeze_output=True, fgraph_name="numba_ofg", **kwargs
     )
 
     if fgraph_cache_key is None:

{pytensor-2.36.2 → pytensor-2.37.0}/pytensor/link/numba/dispatch/elemwise.py

@@ -14,11 +14,11 @@ from pytensor.link.numba.cache import (
 )
 from pytensor.link.numba.dispatch import basic as numba_basic
 from pytensor.link.numba.dispatch.basic import (
-    create_tuple_string,
     numba_funcify_and_cache_key,
     register_funcify_and_cache_key,
     register_funcify_default_op_cache_key,
 )
+from pytensor.link.numba.dispatch.string_codegen import create_tuple_string
 from pytensor.link.numba.dispatch.vectorize_codegen import (
     _vectorized,
     encode_literals,
@@ -365,6 +365,7 @@ def numba_funcify_Elemwise(op, node, **kwargs):
                 output_bc_patterns_enc,
                 output_dtypes_enc,
                 inplace_pattern_enc,
+                True,  # allow_core_scalar
                 (),  # constant_inputs
                 inputs,
                 core_output_shapes,  # core_shapes
@@ -466,6 +467,16 @@ def numba_funcify_DimShuffle(op: DimShuffle, node, **kwargs):
 
         return squeeze_to_0d
 
+    elif op.input_ndim == 0:
+        # DimShuffle can only be an expand_dims or a no_op
+        # This branch uses asarray in case we get a scalar due to https://github.com/numba/numba/issues/10358
+        new_shape = shape_template
+        new_strides = strides_template
+
+        @numba_basic.numba_njit
+        def dimshuffle(x):
+            return as_strided(np.asarray(x), shape=new_shape, strides=new_strides)
+
     else:
 
         @numba_basic.numba_njit
@@ -490,7 +501,7 @@ def numba_funcify_DimShuffle(op: DimShuffle, node, **kwargs):
 
             return as_strided(x, shape=new_shape, strides=new_strides)
 
-    cache_version = 1
+    cache_version = 2
     return dimshuffle, cache_version