feat(mlx): pt.random support with mlx backend

[williambdean]

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

• New feature / enhancement
• Bug fix
• Documentation
• Maintenance
• Other (please specify):

[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]

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

[jessegrabowski]

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

[ricardoV94]

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

[ricardoV94]

MvNormal supports different decomposition strategies, you may want to implement like numba dispatch/op.perform which is more low level if mx.random.multivariate_normal doesn't support them. Or if it's unfeasible issue a warning that it isn't respected and will fallback to svd (if it wasn't svd to begin with)

[ricardoV94]

raise at dispatch time already