jax.lax.broadcasted_iota

jax.lax.broadcasted_iota(dtype, shape, dimension)

Convenience wrapper around iota.

Parameters:

• dtype (jax.typing.DTypeLike)

• shape (Sequence[int | Any])

• dimension (int)

Return type:

Array