handle rng input

Author: williambdean

pytensor/link/mlx/dispatch/__init__.py

@@ -14,4 +14,5 @@
 import pytensor.link.mlx.dispatch.sort
 import pytensor.link.mlx.dispatch.slinalg
 import pytensor.link.mlx.dispatch.nlinalg
+import pytensor.link.mlx.dispatch.random
 # isort: on

pytensor/link/mlx/dispatch/core.py

@@ -300,6 +300,16 @@ def _coerce_to_int(value):
     )
 
 
+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/link/mlx/linker.py

@@ -69,9 +69,16 @@ def create_thunk_inputs(self, storage_map):
         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