feat(mlx): pt.random support with mlx backend by williambdean · Pull Request #1979 · pymc-devs/pytensor

williambdean · 2026-03-16T18:19:39Z

Description

Basic support for mlx random generation.

import pytensor.tensor as pt

data = pt.random.normal(size=(52, 3))

# Works now
data.eval(mode="MLX")

They have limited support. Missing Gamma distribution. Could support additional ones
with basic transformations. i.e. pt.abs(pt.random.normal(...)) ~ Half Normal

MLX Reference: https://ml-explore.github.io/mlx/build/html/python/random.html

Related Issue

Closes #
Related to #

Checklist

Checked that the pre-commit linting/style checks pass
Included tests that prove the fix is effective or that the new feature works
Added necessary documentation (docstrings and/or example notebooks)
If you are a pro: each commit corresponds to a relevant logical change

Type of change

ricardoV94

missing rng outputs /updates (so consecutive calls get updated rng)

There should be tests in numba/jax you can use as template. Jax is going to be more similar

ricardoV94 · 2026-03-16T20:08:51Z

        thunk_inputs = []
        for n in self.fgraph.inputs:
            sinput = storage_map[n]
+            if isinstance(sinput[0], Generator):


you need to do the same dance jax linker does with shared Generator variables

jessegrabowski · 2026-03-30T11:40:54Z

#2010 caused conflicts for this PR. You will need to rebase.

ricardoV94 · 2026-04-01T16:42:43Z

+    def sample_fn(rng_key, size, dtype, p):
+        p = mx.array(p)
+        if size is None:
+            shape = p.shape


you always need the shape? You didn't need it in the categorical. I would assume you only need when one of the parameters doesn't go in the random function. If so that would take a lot of boilerplate away from your dispatches

my comment wasn't about Bernoulli specifically, I would expect you don't need to define shape explicitly (when the user didn't do it themselves) most of the time

jessegrabowski · 2026-04-22T14:14:57Z

why?

williambdean · 2026-04-22T18:06:34Z

Accident

jessegrabowski · 2026-06-18T20:37:35Z

@williambdean what enticements are required to have you get this over the line? It seems to be the last major blocker to mlx-only workflows in pymc

williambdean · 2026-06-18T20:45:28Z

I will take look tonight and tomororw

williambdean · 2026-06-19T03:48:19Z

Changes since last review:

Simplified Normal shape boilerplate by delegating to mx.random.normal(loc, scale) when size is None
Added LogNormal, HalfNormal, Exponential, Logistic, Cauchy dispatches via transform-based sampling

Remaining distributions (Gamma, Poisson, StudentT, Beta, etc.) are more involved and can be discussed as follow-up.

williambdean · 2026-06-19T15:54:54Z

Whats the deal with the benchmark?

williambdean · 2026-06-22T19:50:50Z

Can you take look when you get the time @jessegrabowski 🙏

jessegrabowski · 2026-06-22T20:03:02Z

Yes and I recognize the hypocrisy of asking you to do it then not providing a review :)

re: benchmark, there's something about we already used 256 slots or something. @ricardoV94 knows what is wrong but can't be bothered to fix it.

williambdean · 2026-06-22T20:50:42Z

ha, all good 😄

cetagostini · 2026-06-28T16:13:58Z

Hey @williambdean 👋 — went through this and found a few things while testing the dispatches, so I pushed a commit on top (515092d1) addressing them. Hope you don't mind!

HalfNormal was wrong for nonzero loc: it sampled abs(loc + scale*z), but the support must be [loc, ∞), i.e. loc + scale*|z| (verified against scipy.stats.halfnorm).
MvNormal: batched params + an explicit size raised a reshape error, and an empty batch dim (size=(0,)) actually segfaulted the MLX compiled matmul. Replaced the two branches with a single broadcasting matmul mean + (A @ z[..., None])[..., 0] and guarded the empty case.
Integers: bounds were cast to the output dtype before sampling, so e.g. dtype="uint8" collapsed [250, 300) to a single value. Now we sample at int64 width and cast the result (matches numpy's sample-then-cast).
Bernoulli/Categorical: cast the bool/uint32 draw back to the declared int dtype.
Boilerplate: folded Normal/Laplace/Gumbel into one loc-scale dispatch (getattr(mx.random, op.name)) + a _shape_from_size helper — this is the shape-boilerplate cleanup @ricardoV94 mentioned.
Non-PCG64 generators now raise a clear NotImplementedError instead of an opaque KeyError.

Added regression tests for each. Branch merges cleanly with main (no conflicts). Did this with AI assistance (Cursor) — happy to tweak anything.

cetagostini

looks good to me!

jessegrabowski · 2026-06-28T16:53:05Z

did you just add a bunch of stuff and then approve your own code with a "lgtm" comment?

williambdean · 2026-06-28T17:05:00Z

The vibes look good

jessegrabowski · 2026-06-28T17:07:35Z

I'll give a review today, sorry for being slow

cetagostini · 2026-06-28T17:25:35Z

did you just add a bunch of stuff and then approve your own code with a "lgtm" comment?

kinda 😅 I saw a few minor issues (only 2 files to adjust), and solve them, everything was almost there. But ofc, let's wait your reviews!

williambdean · 2026-07-27T14:16:01Z

What is needed to get this over the finish line?

jessegrabowski

Looks good on my end, but I made some small suggestions on the tests. I say we put it in and see what happens. Highest risk for correctness is in fgraph_convert, I'd ask that you make sure that's carefully tested -- check that you're doing at least as much as the jax suite is doing.

jessegrabowski · 2026-07-27T14:29:07Z

+        pytensor.function([], rv, mode="MLX", updates=srng.updates())
+
+
+def test_beta_not_implemented():


These "not implemented" tests can be removed

jessegrabowski · 2026-07-27T14:32:30Z

+    )
+
+
+def test_gumbel_shape_dtype():


All of the test_shape_dtype tests could be collapsed to a single parameterized test, would be a bit easier to maintain in the long term

jessegrabowski

Rebased this onto main in a worktree — 12 commits, clean, no conflicts. With that on top it trains a convolutional VAE on MLX end to end (reparameterization trick and all), 30 epochs in 1.7s against numba's 23.8s for the same final loss. The AttributeError: 'int' object has no attribute 'astype' I hit before the rebase was version skew, not this PR.

Four small comments below, none blocking. Checked the MvNormal factorizations numerically — A Aᵀ == cov for all three methods.

jessegrabowski · 2026-08-22T04:53:56Z

+def mlx_sample_fn_loc_scale(op, node):
+    """Loc-scale families: MLX names the standard sampler like the Op, so draw
+    it and apply ``loc + scale * z`` (mirrors the JAX dispatch)."""
+    mlx_op = getattr(mx.random, op.name)


Looking the MLX function up by op.name ties PyTensor's Op names to MLX's API surface. Works today, but a rename either side turns into an AttributeError here instead of the NotImplementedError the singledispatch fallback already gives you. An explicit {ptr.NormalRV: mx.random.normal, ...} is the same length and fails predictably.

jessegrabowski · 2026-08-22T04:53:56Z

+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


Everything else goes through _shape_from_size, which broadcasts the parameters when size is None; bernoulli and categorical pass None and let MLX derive it from p. If that's deliberate a one-liner saying so would help — otherwise it reads as drift from the helper. Same on line 130.

jessegrabowski · 2026-08-22T04:53:56Z

+
+        # Factor ``cov = A @ A.T`` so that ``mean + A @ z`` has covariance ``cov``.
+        if method == "cholesky":
+            A = mx.linalg.cholesky(cov, stream=mx.cpu)


Worth a word on why these three run on the CPU stream — I assume MLX's GPU linalg doesn't cover cholesky/svd/eigh, but the next reader has to guess.

jessegrabowski · 2026-08-22T04:53:56Z

+        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)
+        return -scale * mx.log(u)


mx.random.uniform is [low, high), so u = 0 gives -inf here and at the logistic on 233. I couldn't produce one in 200M draws — smallest u was 7.9e-9 — so this is insurance rather than a fix, but -scale * mx.log1p(-u) costs nothing.

williambdean · 2026-08-24T05:39:43Z

Thanks for the feedback, @jessegrabowski . Should be all addressed!

jessegrabowski · 2026-08-24T12:51:04Z

tests are failing :(

jessegrabowski · 2026-08-27T17:31:01Z

@williambdean i think some of these test failures aren't you and were solved by #2376 . A simple rebase will clear many/most/all.

…ethods, permutation error at dispatch time

- LogNormal: exp(normal(mu, sigma)) - HalfNormal: abs(normal(loc, scale)) - Exponential: -scale * log(uniform) (inverse CDF) - Logistic: loc + scale * log(u / (1 - u)) (inverse CDF) - Cauchy: loc + scale * tan(pi * (u - 0.5)) (inverse CDF) - Normal: pass loc/scale directly to MX when size is None - Add tests for all new distributions

- Consolidate Normal/Laplace/Gumbel into a single loc-scale dispatch via getattr(mx.random, op.name) and add a _shape_from_size helper, removing the per-distribution shape/affine boilerplate (mirrors the JAX dispatch). - HalfNormal: sample loc + scale*|z| instead of abs(loc + scale*z); the latter was wrong for nonzero loc (support must be [loc, inf)). - MvNormal: replace the reshape-based size branch with a single broadcasting matmul (mean + (A @ z[..., None])[..., 0]) so batched parameters combined with an explicit size no longer crash; guard empty batch dims, which segfaulted the MLX compiled matmul path. - Integers: sample at full int64 width then cast the result, so narrow output dtypes wrap (matching numpy) and wide bounds aren't clipped to MLX's default int32 sampling dtype. - Bernoulli/Categorical: cast the bool/uint32 draw to the declared int dtype. - Raise a clear NotImplementedError for non-PCG64 bit generators instead of an opaque KeyError. Adds regression tests for each fix. Developed with AI assistance (Cursor). Co-authored-by: Cursor <cursoragent@cursor.com>

- _shape_from_size: return a concrete tuple from mx.broadcast_shapes so mypy no longer reports no-any-return. - Tighten docstrings/comments to one-line intent (why, not what). Co-authored-by: Cursor <cursoragent@cursor.com>

mx.random primitives reject float64 regardless of device, while upstream convert_dtype_to_mlx now preserves float64 on CPU. Sample in float32 and widen the draw centrally; use float dtype family in remaining test assertions.

williambdean · 2026-09-03T12:55:49Z

Rebase should be done!

williambdean added the mlx label Mar 16, 2026

williambdean force-pushed the mlx-random branch 2 times, most recently from b620bb1 to d602268 Compare March 16, 2026 18:25

ricardoV94 reviewed Mar 16, 2026

View reviewed changes

Comment thread pytensor/link/mlx/dispatch/random.py Outdated

williambdean force-pushed the mlx-random branch from e6f7371 to 0b4fb85 Compare March 31, 2026 21:14

ricardoV94 reviewed Apr 1, 2026

View reviewed changes

williambdean force-pushed the mlx-random branch from 0af680e to 5d450de Compare April 10, 2026 19:31

williambdean closed this Apr 22, 2026

williambdean deleted the mlx-random branch April 22, 2026 14:14

williambdean restored the mlx-random branch April 22, 2026 18:48

williambdean reopened this Apr 22, 2026

williambdean force-pushed the mlx-random branch from 5d450de to 1b058be Compare June 19, 2026 03:15

cetagostini approved these changes Jun 28, 2026

View reviewed changes

cetagostini requested a review from ricardoV94 June 28, 2026 16:33

jessegrabowski approved these changes Jul 27, 2026

View reviewed changes

jessegrabowski reviewed Aug 22, 2026

View reviewed changes

williambdean and others added 15 commits August 27, 2026 15:45

add currently supported random module

648411c

add test suite

ac3b831

handle rng input

579941d

address review: shared rng dance in linker, xor-fold 128-bit pcg64 seed

9a58850

fix import after core.py renamed to tensor_basic.py

2688290

address review: bernoulli shape boilerplate, mvnormal decomposition m…

0a93349

…ethods, permutation error at dispatch time

fix ruff isort: declare pytensor and tests as known-first-party

c2c4a31

fix import: pytensor.compile.function removed in v3 refactor

bd62f61

leverage shape

f069f62

fix mypy no-any-return and trim over-explanatory comments

eaf7d33

- _shape_from_size: return a concrete tuple from mx.broadcast_shapes so mypy no longer reports no-any-return. - Tighten docstrings/comments to one-line intent (why, not what). Co-authored-by: Cursor <cursoragent@cursor.com>

address review: consolidate loc-scale dispatch, parametrize shape tests

5ee96d6

fix: use float dtype family in test assertions for floatX compat

c783bd4

sample MLX randoms in float32 and cast to output dtype

d6f1884

mx.random primitives reject float64 regardless of device, while upstream convert_dtype_to_mlx now preserves float64 on CPU. Sample in float32 and widen the draw centrally; use float dtype family in remaining test assertions.

williambdean force-pushed the mlx-random branch from 2c5a2ba to d6f1884 Compare August 28, 2026 19:48

		pytensor.function([], rv, mode="MLX", updates=srng.updates())


		def test_beta_not_implemented():

Uh oh!

Conversation

williambdean commented Mar 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Related Issue

Checklist

Type of change

Uh oh!

ricardoV94 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

jessegrabowski commented Mar 30, 2026

Uh oh!

ricardoV94 Apr 1, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

jessegrabowski commented Apr 22, 2026

Uh oh!

williambdean commented Apr 22, 2026

Uh oh!

jessegrabowski commented Jun 18, 2026

Uh oh!

williambdean commented Jun 18, 2026

Uh oh!

williambdean commented Jun 19, 2026

Uh oh!

williambdean commented Jun 19, 2026

Uh oh!

williambdean commented Jun 22, 2026

Uh oh!

jessegrabowski commented Jun 22, 2026

Uh oh!

williambdean commented Jun 22, 2026

Uh oh!

cetagostini commented Jun 28, 2026

Uh oh!

cetagostini left a comment

Choose a reason for hiding this comment

Uh oh!

jessegrabowski commented Jun 28, 2026

Uh oh!

williambdean commented Jun 28, 2026

Uh oh!

jessegrabowski commented Jun 28, 2026

Uh oh!

cetagostini commented Jun 28, 2026

Uh oh!

williambdean commented Jul 27, 2026

Uh oh!

jessegrabowski left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jessegrabowski left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

williambdean commented Aug 24, 2026

Uh oh!

jessegrabowski commented Aug 24, 2026

Uh oh!

williambdean commented Mar 16, 2026 •

edited

Loading

ricardoV94 Apr 1, 2026 •

edited

Loading