jax.lax.iota# jax.lax.iota(dtype, size)[source]# Wraps XLA’s Iota operator. Parameters: dtype (jax.typing.DTypeLike) size (int) Return type: Array
