pytensor 2.36.1__tar.gz → 2.36.3__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pytensor
-Version: 2.36.1
+Version: 2.36.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

{pytensor-2.36.1 → pytensor-2.36.3}/pytensor/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2025-12-21T14:56:52+0000",
+ "date": "2026-01-07T19:04:45+0100",
  "dirty": false,
  "error": null,
- "full-revisionid": "409a4b2a02e64acc0ca12cd0e046e99fd6f535b9",
- "version": "2.36.1"
+ "full-revisionid": "ac11da625fc3155790bc0b4794807232b73de0b8",
+ "version": "2.36.3"
 }
 '''  # END VERSION_JSON
 

{pytensor-2.36.1 → pytensor-2.36.3}/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.3}/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.1 → pytensor-2.36.3}/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.1 → pytensor-2.36.3}/pytensor/gradient.py

@@ -2104,16 +2104,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.1 → pytensor-2.36.3}/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.3}/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.1 → pytensor-2.36.3}/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.1 → pytensor-2.36.3}/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.1 → pytensor-2.36.3}/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.1 → pytensor-2.36.3}/pytensor/link/numba/dispatch/basic.py

@@ -171,40 +171,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):
@@ -478,7 +444,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.3}/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.3}/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,
@@ -466,6 +466,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 +500,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