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