Skip to main content
Ctrl+K

jax.lax.broadcasted_iota#

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

Convenience wrapper around iota.

Parameters:
  • dtype (Union[str, type[Any], dtype, SupportsDType]) –

  • shape (Sequence[Union[int, Any]]) –

  • dimension (int) –

Return type:

Array

previous

jax.lax.broadcast_to_rank

next

jax.lax.cbrt