jax.lax.broadcasted_iota

Warning

This page was created from a pull request (#9655).

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

previous

jax.lax.broadcast

next

jax.lax.broadcast_in_dim