pytensor 3.3.1__tar.gz → 3.3.2__tar.gz

{pytensor-3.3.1/pytensor.egg-info → pytensor-3.3.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pytensor
-Version: 3.3.1
+Version: 3.3.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-3.3.1 → pytensor-3.3.2}/pyproject.toml

@@ -151,6 +151,7 @@ unfixable = [
 
 [tool.ruff.lint.isort]
 lines-after-imports = 2
+known-first-party = ["pytensor", "tests"]
 
 [tool.ruff.lint.per-file-ignores]
 # TODO: Get rid of these:

{pytensor-3.3.1 → pytensor-3.3.2}/pytensor/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2026-09-07T16:04:36+0200",
+ "date": "2026-09-14T16:15:54+0200",
  "dirty": false,
  "error": null,
- "full-revisionid": "bb5967ab72cfc29ee9f47c4a2fd0aaeeb1695268",
- "version": "3.3.1"
+ "full-revisionid": "895f7795e2ae8ea16e3a4097ce4b92c61c8faf76",
+ "version": "3.3.2"
 }
 '''  # END VERSION_JSON
 

{pytensor-3.3.1 → pytensor-3.3.2}/pytensor/d3viz/d3viz.py

@@ -18,7 +18,7 @@ def replace_patterns(x, replace):
 
     Parameters
     ----------
-    s : str
+    x : str
         String on which function is applied
     replace : dict
         `key`, `value` pairs where key is a regular expression and `value` a

{pytensor-3.3.1 → pytensor-3.3.2}/pytensor/graph/fg.py

@@ -321,7 +321,7 @@ class FunctionGraph(AbstractFunctionGraph):
 
         Parameters
         ----------
-        variable : pytensor.graph.basic.Variable
+        var : pytensor.graph.basic.Variable
             The variable to be imported.
         reason : str
             The name of the optimization or operation in progress.

{pytensor-3.3.1 → pytensor-3.3.2}/pytensor/graph/traversal.py

@@ -666,9 +666,9 @@ def toposort_with_orderings(
     Parameters
     ----------
     graphs : list or tuple of Variable instances
-        Graph inputs.
-    outputs : list or tuple of Apply instances
         Graph outputs.
+    blockers : list or tuple of Variable instances
+        Graph inputs.
     orderings : dict
         Keys are `Apply` or `Variable` instances, values are lists of `Apply` or `Variable` instances.
 

{pytensor-3.3.1 → pytensor-3.3.2}/pytensor/link/mlx/dispatch/__init__.py

@@ -17,4 +17,5 @@ import pytensor.link.mlx.dispatch.extra_ops
 import pytensor.link.mlx.dispatch.pad
 import pytensor.link.mlx.dispatch.sort
 import pytensor.link.mlx.dispatch.linalg
+import pytensor.link.mlx.dispatch.random
 # isort: on

{pytensor-3.3.1 → pytensor-3.3.2}/pytensor/link/mlx/dispatch/basic.py

@@ -132,7 +132,9 @@ def mlx_typify_tensor(data, dtype=None, **kwargs):
     # and it does so on the CPU too, where float64 is perfectly usable
     if dtype is None and data.dtype == np.float64 and float64_supported():
         dtype = mx.float64
-    return _nan_safe_constant(data, dtype=dtype)
+    # MLX's elementwise kernels misread a non-contiguous buffer, and rewriting
+    # produces such arrays -- `triu` of a transpose, for one
+    return _nan_safe_constant(np.asarray(data, order="C"), dtype=dtype)
 
 
 @mlx_typify.register(slice)

{pytensor-3.3.1 → pytensor-3.3.2}/pytensor/link/mlx/dispatch/linalg/solvers.py

@@ -33,12 +33,22 @@ def mlx_funcify_Solve(op, node, **kwargs):
 @mlx_funcify.register(SolveTriangular)
 def mlx_funcify_SolveTriangular(op, node, **kwargs):
     lower = op.lower
+    unit_diagonal = op.unit_diagonal
     A_dtype = getattr(mx, node.inputs[0].dtype)
     b_dtype = getattr(mx, node.inputs[1].dtype)
 
     def solve_triangular(A, b):
+        A = A.astype(stream=mx.cpu, dtype=A_dtype)
+
+        if unit_diagonal:
+            # MLX's `solve_triangular` has no `unit_diagonal`. LAPACK's `trtrs`
+            # never reads the diagonal in that mode, so overwriting it with ones
+            # gives the same answer.
+            diagonal_mask = mx.eye(A.shape[-1], dtype=mx.bool_, stream=mx.cpu)
+            A = mx.where(diagonal_mask, mx.array(1, dtype=A_dtype), A, stream=mx.cpu)
+
         return mx.linalg.solve_triangular(
-            A.astype(stream=mx.cpu, dtype=A_dtype),
+            A,
             b.astype(stream=mx.cpu, dtype=b_dtype),
             upper=not lower,
             stream=mx.cpu,

pytensor-3.3.2/pytensor/link/mlx/dispatch/random.py

@@ -0,0 +1,276 @@
+from functools import singledispatch
+
+import mlx.core as mx
+from numpy.random import Generator
+
+import pytensor.tensor.random.basic as ptr
+from pytensor.link.mlx.dispatch.basic import mlx_funcify, mlx_typify
+from pytensor.link.mlx.dispatch.tensor_basic import (
+    convert_dtype_to_mlx,
+    mlx_to_list_shape,
+)
+
+
+def numpy_generator_to_mlx_key(rng: Generator) -> mx.array:
+    """Convert a NumPy Generator to an MLX random key.
+
+    MLX keys are 64-bit, so we XOR-fold the two halves of the 128-bit PCG64
+    state to keep all of its entropy.
+    """
+    state = rng.bit_generator.state
+    if state["bit_generator"] not in ("PCG64", "PCG64DXSM"):
+        raise NotImplementedError(
+            "MLX RNG conversion only supports the PCG64 bit generator, got "
+            f"{state['bit_generator']}."
+        )
+    state_128 = int(state["state"]["state"])
+    upper = (state_128 >> 64) & 0xFFFFFFFFFFFFFFFF
+    lower = state_128 & 0xFFFFFFFFFFFFFFFF
+    return mx.random.key(upper ^ lower)
+
+
+def _shape_from_size(size, *parameters) -> list[int] | tuple[int, ...]:
+    """Sampling shape: ``size`` if given, else the broadcast of the parameters."""
+    if size is not None:
+        return mlx_to_list_shape(size)
+    return tuple(mx.broadcast_shapes(*(p.shape for p in parameters)))
+
+
+@mlx_typify.register(Generator)
+def mlx_typify_Generator(rng, **kwargs):
+    return numpy_generator_to_mlx_key(rng)
+
+
+@mlx_funcify.register(ptr.RandomVariable)
+def mlx_funcify_RandomVariable(op, node, **kwargs):
+    rv = node.outputs[1]
+    out_dtype = rv.type.dtype
+
+    # MLX random primitives reject float64 regardless of device, so sample in
+    # float32 and widen the draw to the declared output dtype afterwards.
+    mlx_out_dtype = convert_dtype_to_mlx(out_dtype)
+    sample_dtype = mx.float32 if mlx_out_dtype == mx.float64 else mlx_out_dtype
+
+    sample_fn_inner = mlx_sample_fn(op, node)
+
+    def sample_fn(rng, size, *parameters):
+        new_rng, sampling_key = mx.random.split(rng, num=2)
+        sample = sample_fn_inner(sampling_key, size, sample_dtype, *parameters)
+        if sample.dtype != mlx_out_dtype:
+            sample = sample.astype(mlx_out_dtype)
+        return (new_rng, sample)
+
+    return sample_fn
+
+
+@singledispatch
+def mlx_sample_fn(op, node):
+    raise NotImplementedError(
+        f"No MLX implementation for the given distribution: {op.name}"
+    )
+
+
+@mlx_sample_fn.register(ptr.NormalRV)
+def mlx_sample_fn_normal(op, node):
+    def sample_fn(rng_key, size, dtype, loc, scale):
+        mlx_dtype = convert_dtype_to_mlx(dtype)
+        loc = mx.array(loc, dtype=mlx_dtype)
+        scale = mx.array(scale, dtype=mlx_dtype)
+        shape = _shape_from_size(size, loc, scale)
+        return loc + scale * mx.random.normal(shape=shape, dtype=mlx_dtype, key=rng_key)
+
+    return sample_fn
+
+
+@mlx_sample_fn.register(ptr.LaplaceRV)
+def mlx_sample_fn_laplace(op, node):
+    def sample_fn(rng_key, size, dtype, loc, scale):
+        mlx_dtype = convert_dtype_to_mlx(dtype)
+        loc = mx.array(loc, dtype=mlx_dtype)
+        scale = mx.array(scale, dtype=mlx_dtype)
+        shape = _shape_from_size(size, loc, scale)
+        return loc + scale * mx.random.laplace(
+            shape=shape, dtype=mlx_dtype, key=rng_key
+        )
+
+    return sample_fn
+
+
+@mlx_sample_fn.register(ptr.GumbelRV)
+def mlx_sample_fn_gumbel(op, node):
+    def sample_fn(rng_key, size, dtype, loc, scale):
+        mlx_dtype = convert_dtype_to_mlx(dtype)
+        loc = mx.array(loc, dtype=mlx_dtype)
+        scale = mx.array(scale, dtype=mlx_dtype)
+        shape = _shape_from_size(size, loc, scale)
+        return loc + scale * mx.random.gumbel(shape=shape, dtype=mlx_dtype, key=rng_key)
+
+    return sample_fn
+
+
+@mlx_sample_fn.register(ptr.UniformRV)
+def mlx_sample_fn_uniform(op, node):
+    def sample_fn(rng_key, size, dtype, low, high):
+        mlx_dtype = convert_dtype_to_mlx(dtype)
+        low = mx.array(low, dtype=mlx_dtype)
+        high = mx.array(high, dtype=mlx_dtype)
+        shape = _shape_from_size(size, low, high)
+        return mx.random.uniform(
+            low=low, high=high, shape=shape, dtype=mlx_dtype, key=rng_key
+        )
+
+    return sample_fn
+
+
+@mlx_sample_fn.register(ptr.IntegersRV)
+def mlx_sample_fn_integers(op, node):
+    def sample_fn(rng_key, size, dtype, low, high):
+        low = mx.array(low)
+        high = mx.array(high)
+        shape = _shape_from_size(size, low, high)
+        # Sample at full int64 width and cast the result: PyTensor casts the
+        # output, not the bounds, so narrow/wide dtypes don't corrupt the range.
+        return mx.random.randint(
+            low=low, high=high, shape=shape, dtype=mx.int64, key=rng_key
+        ).astype(convert_dtype_to_mlx(dtype))
+
+    return sample_fn
+
+
+@mlx_sample_fn.register(ptr.BernoulliRV)
+def mlx_sample_fn_bernoulli(op, node):
+    def sample_fn(rng_key, size, dtype, p):
+        p = mx.array(p)
+        shape = mlx_to_list_shape(size) if size is not None else None
+        # MLX draws bool; PyTensor declares an int dtype.
+        return mx.random.bernoulli(p=p, shape=shape, key=rng_key).astype(
+            convert_dtype_to_mlx(dtype)
+        )
+
+    return sample_fn
+
+
+@mlx_sample_fn.register(ptr.CategoricalRV)
+def mlx_sample_fn_categorical(op, node):
+    def sample_fn(rng_key, size, dtype, p):
+        logits = mx.log(mx.array(p))
+        shape = mlx_to_list_shape(size) if size is not None else None
+        # MLX draws uint32; PyTensor declares an int dtype.
+        return mx.random.categorical(
+            logits=logits, axis=-1, shape=shape, key=rng_key
+        ).astype(convert_dtype_to_mlx(dtype))
+
+    return sample_fn
+
+
+@mlx_sample_fn.register(ptr.MvNormalRV)
+def mlx_sample_fn_mvnormal(op, node):
+    method = op.method
+
+    def sample_fn(rng_key, size, dtype, mean, cov):
+        mlx_dtype = convert_dtype_to_mlx(dtype)
+        mean = mx.array(mean, dtype=mlx_dtype)
+        cov = mx.array(cov, dtype=mlx_dtype)
+
+        n = cov.shape[-1]
+        if size is not None:
+            batch_shape = mlx_to_list_shape(size)
+        else:
+            batch_shape = mx.broadcast_shapes(mean.shape[:-1], cov.shape[:-2])
+
+        if 0 in tuple(batch_shape):
+            # Empty batch dim crashes MLX's compiled matmul; the draw is empty anyway.
+            return mx.broadcast_to(mean, [*batch_shape, n])
+
+        # Factor ``cov = A @ A.T`` so that ``mean + A @ z`` has covariance ``cov``.
+        if method == "cholesky":
+            A = mx.linalg.cholesky(cov, stream=mx.cpu)
+        elif method == "svd":
+            U, s, _ = mx.linalg.svd(cov, stream=mx.cpu)
+            A = U * mx.sqrt(s)[..., None, :]
+        else:  # eigh
+            w, vecs = mx.linalg.eigh(cov, stream=mx.cpu)
+            A = vecs * mx.sqrt(w)[..., None, :]
+
+        z = mx.random.normal(shape=[*batch_shape, n], dtype=mlx_dtype, key=rng_key)
+        return mean + (A @ z[..., None])[..., 0]
+
+    return sample_fn
+
+
+@mlx_sample_fn.register(ptr.PermutationRV)
+def mlx_sample_fn_permutation(op, node):
+    if op.batch_ndim(node):
+        raise NotImplementedError(
+            "MLX random.permutation does not support batch dimensions."
+        )
+
+    def sample_fn(rng_key, size, dtype, x):
+        return mx.random.permutation(x, key=rng_key)
+
+    return sample_fn
+
+
+@mlx_sample_fn.register(ptr.LogNormalRV)
+def mlx_sample_fn_lognormal(op, node):
+    def sample_fn(rng_key, size, dtype, mu, sigma):
+        mlx_dtype = convert_dtype_to_mlx(dtype)
+        mu = mx.array(mu, dtype=mlx_dtype)
+        sigma = mx.array(sigma, dtype=mlx_dtype)
+        shape = _shape_from_size(size, mu, sigma)
+        z = mx.random.normal(shape=shape, dtype=mlx_dtype, key=rng_key)
+        return mx.exp(mu + sigma * z)
+
+    return sample_fn
+
+
+@mlx_sample_fn.register(ptr.HalfNormalRV)
+def mlx_sample_fn_halfnormal(op, node):
+    def sample_fn(rng_key, size, dtype, loc, scale):
+        mlx_dtype = convert_dtype_to_mlx(dtype)
+        loc = mx.array(loc, dtype=mlx_dtype)
+        scale = mx.array(scale, dtype=mlx_dtype)
+        shape = _shape_from_size(size, loc, scale)
+        z = mx.random.normal(shape=shape, dtype=mlx_dtype, key=rng_key)
+        return loc + scale * mx.abs(z)
+
+    return sample_fn
+
+
+@mlx_sample_fn.register(ptr.ExponentialRV)
+def mlx_sample_fn_exponential(op, node):
+    def sample_fn(rng_key, size, dtype, scale):
+        mlx_dtype = convert_dtype_to_mlx(dtype)
+        scale = mx.array(scale, dtype=mlx_dtype)
+        shape = _shape_from_size(size, scale)
+        u = mx.random.uniform(shape=shape, dtype=mlx_dtype, key=rng_key)
+        # log1p(-u) avoids -inf when u=0 (u is in [0, 1))
+        return -scale * mx.log1p(-u)
+
+    return sample_fn
+
+
+@mlx_sample_fn.register(ptr.LogisticRV)
+def mlx_sample_fn_logistic(op, node):
+    def sample_fn(rng_key, size, dtype, loc, scale):
+        mlx_dtype = convert_dtype_to_mlx(dtype)
+        loc = mx.array(loc, dtype=mlx_dtype)
+        scale = mx.array(scale, dtype=mlx_dtype)
+        shape = _shape_from_size(size, loc, scale)
+        u = mx.random.uniform(shape=shape, dtype=mlx_dtype, key=rng_key)
+        return loc + scale * mx.log(u / (1 - u))
+
+    return sample_fn
+
+
+@mlx_sample_fn.register(ptr.CauchyRV)
+def mlx_sample_fn_cauchy(op, node):
+    def sample_fn(rng_key, size, dtype, loc, scale):
+        mlx_dtype = convert_dtype_to_mlx(dtype)
+        loc = mx.array(loc, dtype=mlx_dtype)
+        scale = mx.array(scale, dtype=mlx_dtype)
+        shape = _shape_from_size(size, loc, scale)
+        u = mx.random.uniform(shape=shape, dtype=mlx_dtype, key=rng_key)
+        return loc + scale * mx.tan(mx.pi * (u - 0.5))
+
+    return sample_fn

{pytensor-3.3.1 → pytensor-3.3.2}/pytensor/link/mlx/dispatch/tensor_basic.py

@@ -239,6 +239,16 @@ def _coerce_to_int(value):
         raise
 
 
+def mlx_to_list_shape(size) -> list[int]:
+    """Convert a size value (mx.array, np.ndarray, or sequence) to a plain Python list of ints.
+
+    Used by random variable dispatch to normalise the ``size`` argument, which
+    PyTensor may pass as an ``mx.array`` or ``np.ndarray`` rather than a plain
+    Python list.
+    """
+    return [_coerce_to_int(x) for x in size]
+
+
 def _rethrow_dynamic_shape_error(exc):
     msg = str(exc)
     if "[eval] Attempting to eval an array during function transformations" in msg:

pytensor-3.3.2/pytensor/link/mlx/linker.py

@@ -0,0 +1,143 @@
+import warnings
+
+from pytensor.compile.sharedvalue import SharedVariable, shared
+from pytensor.link.basic import JITLinker
+
+
+class MLXLinker(JITLinker):
+    """A `Linker` that JIT-compiles NumPy-based operations using Apple's MLX."""
+
+    required_rewrites = ("minimum_compile",)
+    incompatible_rewrites = (
+        "cxx_only",
+        "BlasOpt",
+        "local_careduce_fusion",
+        "inplace",
+        "scan_reduce_trace_prealloc",
+        "inline_einsum",
+    )
+
+    def __init__(self, use_compile=True, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.gen_functors = []
+        self.use_compile = use_compile
+
+    def fgraph_convert(self, fgraph, input_storage, storage_map, **kwargs):
+        """Convert a PyTensor FunctionGraph to an MLX-compatible function.
+
+        Parameters
+        ----------
+        fgraph : FunctionGraph
+            The function graph to convert
+
+        Returns
+        -------
+        callable
+            An MLX-compatible function
+        """
+        from pytensor.link.mlx.dispatch import mlx_funcify
+        from pytensor.tensor.random.type import RandomType
+
+        shared_rng_inputs = [
+            inp
+            for inp in fgraph.inputs
+            if (isinstance(inp, SharedVariable) and isinstance(inp.type, RandomType))
+        ]
+
+        # Replace any shared RNG inputs so that their values can be updated in place
+        # without affecting the original RNG container. This is necessary because
+        # MLX does not accept Generators as inputs, and they will have to
+        # be typified
+        if shared_rng_inputs:
+            warnings.warn(
+                f"The RandomType SharedVariables {shared_rng_inputs} will not be used "
+                f"in the compiled MLX graph. Instead a copy will be used.",
+                UserWarning,
+            )
+            new_shared_rng_inputs = [
+                shared(inp.get_value(borrow=False)) for inp in shared_rng_inputs
+            ]
+
+            fgraph.replace_all(
+                zip(shared_rng_inputs, new_shared_rng_inputs, strict=True),
+                import_missing=True,
+                reason="MLXLinker.fgraph_convert",
+            )
+
+            for old_inp, new_inp in zip(
+                shared_rng_inputs, new_shared_rng_inputs, strict=True
+            ):
+                new_inp_storage = [new_inp.get_value(borrow=True)]
+                storage_map[new_inp] = new_inp_storage
+                old_inp_storage = storage_map.pop(old_inp)
+                # Find index of old_inp_storage in input_storage
+                for input_storage_idx, input_storage_item in enumerate(input_storage):
+                    # We have to establish equality based on identity because input_storage may contain numpy arrays
+                    if input_storage_item is old_inp_storage:
+                        break
+                else:  # no break
+                    raise ValueError()
+                input_storage[input_storage_idx] = new_inp_storage
+                # We need to change the order of the inputs of the FunctionGraph
+                # so that the new input is in the same position as to old one,
+                # to align with the storage_map. We hope this is safe!
+                old_inp_fgraph_index = fgraph.inputs.index(old_inp)
+                fgraph.remove_input(
+                    old_inp_fgraph_index,
+                    reason="MLXLinker.fgraph_convert",
+                )
+                fgraph.inputs.remove(new_inp)
+                fgraph.inputs.insert(old_inp_fgraph_index, new_inp)
+
+        return mlx_funcify(
+            fgraph,
+            input_storage=input_storage,
+            storage_map=storage_map,
+            **kwargs,
+        )
+
+    def jit_compile(self, fn):
+        import mlx.core as mx
+
+        from pytensor.link.mlx.dispatch import mlx_typify
+
+        if not self.use_compile:
+            # Skip compilation and just return the function with MLX typification
+            def fn_no_compile(*inputs):
+                return fn(*(mlx_typify(inp) for inp in inputs))
+
+            return fn_no_compile
+
+        inner_fn = mx.compile(fn)
+
+        def fn(*inputs, inner_fn=inner_fn):
+            return inner_fn(*(mlx_typify(inp) for inp in inputs))
+
+        return fn
+
+    def create_thunk_inputs(self, storage_map):
+        """Create inputs for the MLX thunk.
+
+        Parameters
+        ----------
+        storage_map : dict
+            Map from variables to their storage
+
+        Returns
+        -------
+        list
+            The inputs for the thunk
+        """
+        from numpy.random import Generator
+
+        from pytensor.link.mlx.dispatch import mlx_typify
+
+        thunk_inputs = []
+        for n in self.fgraph.inputs:
+            sinput = storage_map[n]
+            if isinstance(sinput[0], Generator):
+                # Convert Generator into MLX PRNG key
+                sinput[0] = mlx_typify(sinput[0])
+            thunk_inputs.append(sinput)
+
+        return thunk_inputs

{pytensor-3.3.1 → pytensor-3.3.2}/pytensor/link/numba/dispatch/cython_support.py

@@ -1,8 +1,7 @@
 import importlib
 import re
-from collections.abc import Callable, Mapping
+from collections.abc import Callable
 from dataclasses import dataclass
-from typing import Any, cast
 
 import numba
 import numpy as np
@@ -33,9 +32,7 @@ _C_TO_NUMPY: dict[str, DTypeLike] = {
 @dataclass
 class Signature:
     res_dtype: DTypeLike
-    res_c_type: str
     arg_dtypes: list[DTypeLike]
-    arg_c_types: list[str]
     arg_names: list[str | None]
 
     @property
@@ -92,7 +89,6 @@ class Signature:
 
         arg_dtypes = []
         arg_names: list[str | None] = []
-        arg_c_types = []
         for raw_arg in raw_args.split(b","):
             re_match = re.fullmatch(decl_expr, raw_arg)
             if re_match is None:
@@ -104,7 +100,6 @@ class Signature:
             except KeyError:
                 raise ValueError(f"Unknown C type: {arg_c_type}")
 
-            arg_c_types.append(arg_c_type)
             arg_dtypes.append(arg_dtype)
             name = groups["name"]
             if not name:
@@ -112,36 +107,38 @@ class Signature:
             else:
                 arg_names.append(name.decode())
 
-        return Signature(res_dtype, res_c_type, arg_dtypes, arg_c_types, arg_names)
+        return Signature(res_dtype, arg_dtypes, arg_names)
 
 
-def _available_impls(func: Callable) -> list[tuple[Signature, Any, str]]:
+def _available_impls(func: Callable) -> list[tuple[Signature, str]]:
     """Find all available implementations for a fused cython function.
 
-    Each entry is ``(signature, capsule, capi_name)``, where ``capi_name`` is the key under
+    Each entry is ``(signature, capi_name)``, where ``capi_name`` is the key under
     which the implementation is exported in the module's ``__pyx_capi__`` table. That name is a
     stable, picklable handle for the C function, used to re-resolve its address at runtime.
     """
     impls = []
     mod = importlib.import_module(func.__module__)
 
-    signatures = getattr(func, "__signatures__", None)
-    if signatures is not None:
-        # Cython function with __signatures__ should be fused and thus
-        # indexable
-        func_map = cast(Mapping, func)
-        candidates = [func_map[key] for key in signatures]
+    if getattr(func, "__signatures__", None) is not None:
+        # Cython 3.3 exports typed names and retains the numbered names as aliases.
+        # Python specialization names need not match either C API name.
+        names = [
+            name for name in mod.__pyx_capi__ if name.startswith(f"{func.__name__}[")
+        ]
+        if not names:
+            pattern = re.compile(rf"__pyx_fuse_[0-9_]+{re.escape(func.__name__)}")
+            names = [name for name in mod.__pyx_capi__ if pattern.fullmatch(name)]
     else:
-        candidates = [func]
-    for candidate in candidates:
-        name = candidate.__name__
+        names = [func.__name__]
+    for name in names:
         capsule = mod.__pyx_capi__[name]
         llc = LowLevelCallable(capsule)
         try:
             signature = Signature.from_c_types(llc.signature.encode())
         except KeyError:
             continue
-        impls.append((signature, capsule, name))
+        impls.append((signature, name))
     return impls
 
 
@@ -178,7 +175,7 @@ class _CythonFunctionSpec:
 def wrap_cython_function(func, restype, arg_types):
     impls = _available_impls(func)
     compatible = []
-    for sig, _capsule, capi_name in impls:
+    for sig, capi_name in impls:
         if sig.provides(restype, arg_types):
             compatible.append((sig, capi_name))
 

{pytensor-3.3.1 → pytensor-3.3.2}/pytensor/link/numba/dispatch/linalg/decomposition/qr.py

@@ -34,7 +34,6 @@ def _qr_full_pivot(
         pivoting=pivoting,
         overwrite_a=overwrite_a,
         check_finite=False,
-        lwork=lwork,
     )
 
 
@@ -57,7 +56,6 @@ def _qr_full_no_pivot(
         pivoting=pivoting,
         overwrite_a=overwrite_a,
         check_finite=False,
-        lwork=lwork,
     )
 
 
@@ -80,7 +78,6 @@ def _qr_r_pivot(
         pivoting=pivoting,
         overwrite_a=overwrite_a,
         check_finite=False,
-        lwork=lwork,
     )
 
 
@@ -103,7 +100,6 @@ def _qr_r_no_pivot(
         pivoting=pivoting,
         overwrite_a=overwrite_a,
         check_finite=False,
-        lwork=lwork,
     )
 
 
@@ -126,7 +122,6 @@ def _qr_raw_no_pivot(
         pivoting=pivoting,
         overwrite_a=overwrite_a,
         check_finite=False,
-        lwork=lwork,
     )
 
     return H, tau, R
@@ -151,7 +146,6 @@ def _qr_raw_pivot(
         pivoting=pivoting,
         overwrite_a=overwrite_a,
         check_finite=False,
-        lwork=lwork,
     )
 
     return H, tau, R, P