Control flow in compile/vmap

[ghcollin]

I'm porting some bayesian inferences algorithms to MLX such as Hamiltonian Monte Carlo and other samplers. These algorithms use adaptive/randomly chosen numbers of iterations in their inner loops. As far as I can tell, in MLX, it is only possible to write a while loop using standard python control flow, which requires the array to evaluated, so it can't be compiled. Unfortunately, this is a bit of a blocker because these algorithms also require flattening/unflattening user provided pytrees at every iteration, so uncompiled functions get bogged down in a mess of pure python logistical operations.

Is there a plan to bring a while loop-like construct to MLX? It's super useful for scientific computing, as it unlock all sorts of adaptive algorithms.

[awni]

I think it would be cool to add something like that which could go in the graph. It's not a top priority so it may take some time to get there, but we'd also be open to contributions if you are interested. There has been a feature req open for this for a while now.