jax.experimental.pallas.loop

jax.experimental.pallas.loop(lower, upper, *, step=1, unroll=None)

Returns a decorator that calls the decorated function in a loop.

Parameters:

• lower (Array | ndarray | bool | number | bool | int | float | complex | TypedNdArray)

• upper (Array | ndarray | bool | number | bool | int | float | complex | TypedNdArray)

• step (Array | ndarray | bool | number | bool | int | float | complex | TypedNdArray)

• unroll (int | bool | None)

Return type:

Callable[[Callable[[Array], None]], None]