flax
573cd25d - Use eval_shape for initial call to body

Commit

283 days ago

Use eval_shape for initial call to body

References

#5141 - Handle pure bodies in nnx.fori_loop

Author

samanklesaria

samanklesaria

Parents

Loading