jax.lax.iota

jax.lax.iota(dtype, size)

Wraps XLA’s Iota operator.

Parameters:

• dtype (Union[str, type[Any], dtype, SupportsDType])

• size (int)

Return type:

Array