ForiLoop

Functional loop from lower to upper bound.

Abstract Signature:

ForiLoop(lower: int, upper: int, body_fun: Callable, init_val)

PyTorch

API: —
Strategy: Plugin (transform_for_loop)

JAX (Core)

API: jax.lax.fori_loop
Strategy: Direct Mapping
Official Docs ↗

TensorFlow

API: tf.while_loop
Strategy: Direct Mapping
Official Docs ↗

Flax NNX

API: flax.nnx.fori_loop
Strategy: Direct Mapping
Official Docs ↗

PaxML / Praxis

API: jax.lax.fori_loop
Strategy: Direct Mapping
Official Docs ↗