WhileLoop

Constructs a loop that executes body_fn while cond_fn is true.

Abstract Signature:

WhileLoop(cond_fn: Callable, body_fn: Callable, init_val)

PyTorch

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

JAX (Core)

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

NumPy

API: —
Strategy: Custom / Partial

Keras

API: keras.ops.while_loop
Strategy: Direct Mapping
Official Docs ↗

TensorFlow

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

Apple MLX

API: —
Strategy: Custom / Partial

Flax NNX

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

PaxML / Praxis

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