jax.lax.broadcasted_iota

jax.lax.broadcasted_iota(dtype, shape, dimension, _sharding=None)

Convenience wrapper around iota.

Parameters:

• dtype (DTypeLike)

• shape (Shape)

• dimension (int)

Return type:

Array