pytensor 3.3.0__tar.gz → 3.3.2__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pytensor
-Version: 3.3.0
+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
@@ -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<=0.66.0,>=0.58
+Requires-Dist: numba<=0.67.0,>=0.58
 Requires-Dist: filelock>=3.15
 Provides-Extra: complete
 Requires-Dist: pytensor[jax]; extra == "complete"
@@ -57,7 +57,7 @@ Provides-Extra: jax
 Requires-Dist: jax; extra == "jax"
 Requires-Dist: jaxlib; extra == "jax"
 Provides-Extra: numba
-Requires-Dist: numba<=0.66.0,>=0.58; extra == "numba"
+Requires-Dist: numba<=0.67.0,>=0.58; extra == "numba"
 Requires-Dist: llvmlite; extra == "numba"
 Provides-Extra: kanren
 Requires-Dist: etuples; extra == "kanren"

{pytensor-3.3.0 → pytensor-3.3.2}/doc/library/graph/index.rst

@@ -17,5 +17,6 @@
     replace
     features
     op
+    traversal
     type
     utils

pytensor-3.3.2/doc/library/graph/traversal.rst

@@ -0,0 +1,8 @@
+.. _libdoc_graph_traversal:
+
+================================================
+:mod:`traversal` -- Graph Traversal Utilities
+================================================
+
+.. automodule:: pytensor.graph.traversal
+   :members:

{pytensor-3.3.0 → pytensor-3.3.2}/pyproject.toml

@@ -50,7 +50,7 @@ dependencies = [
     "setuptools>=59.0.0",
     "scipy>=1,<2",
     "numpy>=2.0",
-    "numba>=0.58,<=0.66.0",
+    "numba>=0.58,<=0.67.0",
     "filelock>=3.15",
 ]
 
@@ -78,7 +78,7 @@ tests = [
 ]
 rtd = ["sphinx>=5.1.0,<6", "pygments", "pydot"]
 jax = ["jax", "jaxlib"]
-numba = ["numba>=0.58,<=0.66.0", "llvmlite"]
+numba = ["numba>=0.58,<=0.67.0", "llvmlite"]
 kanren = [
     "etuples",
     "logical-unification",
@@ -131,7 +131,11 @@ filterwarnings =[
 
 [tool.ruff]
 line-length = 88
-exclude = ["pytensor/_version.py"]
+# _patch_list_to_tuple.py is copied verbatim from numba; keep it diffable upstream
+exclude = [
+    "pytensor/_version.py",
+    "pytensor/link/numba/dispatch/_patch_list_to_tuple.py",
+]
 
 [tool.ruff.format]
 docstring-code-format = true
@@ -147,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.0 → pytensor-3.3.2}/pytensor/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2026-08-12T10:16:27+0200",
+ "date": "2026-09-14T16:15:54+0200",
  "dirty": false,
  "error": null,
- "full-revisionid": "aa1bc7772eb11ed8def2814374aeb5da03da590c",
- "version": "3.3.0"
+ "full-revisionid": "895f7795e2ae8ea16e3a4097ce4b92c61c8faf76",
+ "version": "3.3.2"
 }
 '''  # END VERSION_JSON
 

{pytensor-3.3.0 → pytensor-3.3.2}/pytensor/compile/mode.py

@@ -524,7 +524,9 @@ JAX = Mode(
 )
 PYTORCH = Mode(
     PytorchLinker(),
-    RewriteDatabaseQuery(include=["fast_run"]),
+    # `Composite` has no PyTorch dispatch (it has no `nfunc_spec` to look up in torch),
+    # and torch fuses its own kernels anyway.
+    RewriteDatabaseQuery(include=["fast_run"], exclude=["fusion"]),
 )
 
 MLX = Mode(

{pytensor-3.3.0 → 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.0 → pytensor-3.3.2}/pytensor/gradient.py

@@ -1335,8 +1335,8 @@ def _populate_grad_dict(var_to_app_to_idx, grad_dict, wrt, cost_name=None):
                     raise ValueError(
                         f"{node.op} returned the wrong number of gradient terms."
                     )
-            # We can not enforce this, as AdvancedSubtensor1 has an option to
-            # return the sparse grad for optimization reason.
+            # We can not enforce this, as sparse Ops have an option to return a
+            # structured grad for optimization reason.
 
             #            for ig, i in zip(input_grads, inputs):
             #                if (not isinstance(ig.type, (DisconnectedType, NullType)) and

{pytensor-3.3.0 → pytensor-3.3.2}/pytensor/graph/basic.py

@@ -45,13 +45,12 @@ _MOVED_FUNCTIONS = {
     "graph_inputs",
     "explicit_graph_inputs",
     "vars_between",
-    "orhpans_between",
+    "orphans_between",
     "applys_between",
     "apply_depends_on",
     "truncated_graph_inputs",
     "general_toposort",
     "io_toposort",
-    "list_of_nodes",
     "get_var_by_name",
 }
 

{pytensor-3.3.0 → 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.0 → pytensor-3.3.2}/pytensor/graph/replace.py

@@ -123,8 +123,6 @@ def graph_replace(
         Replace mapping
     strict: bool
         Raise an error if some replacements were not used
-    return_unused: bool
-        Return replacements that were not used
 
     Returns
     -------

{pytensor-3.3.0 → pytensor-3.3.2}/pytensor/graph/rewriting/basic.py

@@ -43,7 +43,6 @@ from pytensor.graph.traversal import (
 )
 from pytensor.graph.utils import AssocList, InconsistencyError
 from pytensor.misc.ordered_set import OrderedSet
-from pytensor.utils import flatten
 
 
 _logger = logging.getLogger("pytensor.graph.rewriting.basic")
@@ -705,6 +704,15 @@ class MergeOptimizer(GraphRewriter):
             fgraph.attach_feature(MergeFeature())
 
     def apply(self, fgraph):
+        def has_direct_destroyer(var):
+            """Whether any direct client destroys `var` (view-chain destroyers not seen)."""
+            return any(
+                i in destroyed_idxs
+                for client, i in fgraph.clients[var]
+                if client.op.destroy_map
+                for destroyed_idxs in client.op.destroy_map.values()
+            )
+
         sched = fgraph.merge_feature.scheduled
         nb_fail = 0
         t0 = time.perf_counter()
@@ -755,15 +763,17 @@ class MergeOptimizer(GraphRewriter):
                         continue
 
                     if hasattr(fgraph, "destroy_handler"):
-                        # If both nodes have clients that destroy them, we
-                        # can't merge them.
-                        clients = (
-                            fgraph.clients[pairs[0][0]] + fgraph.clients[pairs[0][1]]
-                        )
-                        if any(
-                            i in flatten(c.op.destroy_map.values())
-                            for c, i in clients
-                            if c.op.destroy_map
+                        # Merging two variables that both have a destroyer is
+                        # always rejected by `DestroyHandler` validation
+                        # ("multiple destroyers") -- skip those outright
+                        # instead of paying a checkpoint + failed validate +
+                        # revert per scheduled pair. A single-destroyer merge
+                        # may still be feasible, and destroyers reachable only
+                        # through view chains are rare: both are left to
+                        # validation, keeping this check O(clients) with no
+                        # droot recomputation.
+                        if has_direct_destroyer(pairs[0][0]) and has_direct_destroyer(
+                            pairs[0][1]
                         ):
                             continue
 

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

@@ -145,12 +145,13 @@ def apply_ancestors(
     graphs: Iterable[Variable],
     blockers: Iterable[Variable] | None = None,
 ) -> Generator[Apply, None, None]:
-    """Return the Apply nodes that contribute to those in given graphs (inclusive)."""
+    """Return the Apply nodes that contribute to those in given graphs, excluding the owners of the blockers."""
     seen = {None}  # This filters out Variables without an owner
+    blockers = frozenset(blockers or ())
     for var in ancestors(graphs, blockers):
         # For multi-output nodes, we'll see multiple variables
         # but we should only yield the Apply once
-        if (apply := var.owner) not in seen:
+        if var not in blockers and (apply := var.owner) not in seen:
             yield apply
             seen.add(apply)
     return
@@ -665,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.0 → 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.0 → 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.0 → pytensor-3.3.2}/pytensor/link/mlx/dispatch/linalg/solvers.py

@@ -4,6 +4,7 @@ import mlx.core as mx
 
 from pytensor.link.mlx.dispatch.basic import mlx_funcify
 from pytensor.tensor.linalg.solvers.general import Solve
+from pytensor.tensor.linalg.solvers.psd import CholeskySolve
 from pytensor.tensor.linalg.solvers.triangular import SolveTriangular
 
 
@@ -32,15 +33,44 @@ 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,
         )
 
     return solve_triangular
+
+
+@mlx_funcify.register(CholeskySolve)
+def mlx_funcify_CholeskySolve(op, node, **kwargs):
+    lower = op.lower
+    c_dtype = getattr(mx, node.inputs[0].dtype)
+    b_dtype = getattr(mx, node.inputs[1].dtype)
+
+    # MLX has no cho_solve, so with A = L L.T we solve L y = b then L.T x = y.
+    def cho_solve(c, b):
+        c = c.astype(stream=mx.cpu, dtype=c_dtype)
+        b = b.astype(stream=mx.cpu, dtype=b_dtype)
+        c_T = mx.swapaxes(c, -1, -2, stream=mx.cpu)
+        L, L_T = (c, c_T) if lower else (c_T, c)
+
+        y = mx.linalg.solve_triangular(L, b, upper=False, stream=mx.cpu)
+        return mx.linalg.solve_triangular(L_T, y, upper=True, stream=mx.cpu)
+
+    return cho_solve

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.2/pytensor/link/mlx/dispatch/scalar/__init__.py

@@ -0,0 +1,10 @@
+from pytensor.link.mlx.dispatch.scalar import (
+    basic,
+    bessel,
+    beta,
+    erf,
+    erfcinv,
+    gamma,
+    gammainc,
+    math,
+)