mx.scan?

[pharringtonp19]

Is there an equivalent of jax.lax.scan in mx?

[awni]

No we don't have that.

[angeloskath]

@pharringtonp19 in Jax you can backprop through the loop (it is just unrolled). The reason people often use this pattern is because it reduces compilation time significantly, especially if the loop is big. Often compilation even fails due to memory or other issues if the graph is too big.

[angeloskath]

On the other hand reading the docs you referenced it seems you are right so ignore the possible noise. I just have seen the scan pattern to reduce compilation time once too many times and thought it applies.

[awni]

🤔 I'm pretty sure you can backprop through a loop in jax:

import jax
import jax.numpy as jnp


def loop_fn(x):
    for _ in range(5):
        x = jnp.sin(x)
    return x

print(jax.grad(loop_fn)(jnp.array(1.0)))

[awni]

If you use a fixed iteration count in the inner loop, I think it's fine. And indeed the only reason I can see to use a scan would be to minimize compile time like @angeloskath said..

[pharringtonp19]

@angeloskath @awni That was my mistake. I should just have presented a mwe instead of incorrectly referencing the jax documents. Thanks for your time!