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