jax.lax.broadcasted_iota¶

jax.lax.broadcasted_iota(dtype, shape, dimension)[source]¶

Convenience wrapper around iota.

Parameters

dtype (Any) –
shape (Sequence[Union[int, Any]]) –
dimension (int) –

Return type