Is it possible to use JIT with MLX?

[kyrollosyanny]

Hi,
I am not sure if this question even makes sense but I know that JAX and Pytorch have a JIT (just in time compilation mode) that runs much faster than regular mode. This is generally done using a jit decorator on top of the function defintion. I'm wondering if this is something that can also be done with MLX? I have some functions that I'll run thousands of times, so JIT might make sense to speed it up.

Thanks

[evanwashere]

https://ml-explore.github.io/mlx/build/html/usage/compile.html

[kyrollosyanny]

That's exactly what I was looking for. Thanks a lot