modify the pytorch jit

Author: williambdean

pytensor/link/mlx/linker.py

@@ -66,21 +66,41 @@ def conversion_func_register(*args, **kwargs):
         )
 
     def jit_compile(self, fn):
-        """JIT compile an MLX function.
+        import mlx.core as mx
 
-        Parameters
-        ----------
-        fn : callable
-            The function to compile
+        from pytensor.link.mlx.dispatch import mlx_typify
 
-        Returns
-        -------
-        callable
-            The compiled function
-        """
-        import mlx.core as mx
+        class wrapper:
+            def __init__(self, fn, gen_functors):
+                self.fn = mx.compile(fn)
+                self.gen_functors = gen_functors.copy()
+
+            def __call__(self, *inputs, **kwargs):
+                import pytensor.link.utils
+
+                # set attrs
+                for n, fn in self.gen_functors:
+                    setattr(pytensor.link.utils, n[1:], fn)
+
+                # MLX doesn't support np.ndarray as input
+                outs = self.fn(*(mlx_typify(inp) for inp in inputs), **kwargs)
+
+                return outs
+
+                # unset attrs
+                for n, _ in self.gen_functors:
+                    if getattr(pytensor.link.utils, n[1:], False):
+                        delattr(pytensor.link.utils, n[1:])
+
+                return tuple(out.cpu().numpy() for out in outs)
+
+            def __del__(self):
+                del self.gen_functors
+
+        inner_fn = wrapper(fn, self.gen_functors)
+        self.gen_functors = []
 
-        return mx.compile(fn)
+        return inner_fn
 
     def create_thunk_inputs(self, storage_map):
         """Create inputs for the MLX thunk.