jax.lax.broadcasted_iota

jax.lax.broadcasted_iota(dtype, shape, dimension)

围绕 iota 的便捷包装器。

参数：

• dtype (DTypeLike)

• shape (Shape)

• dimension (int)

返回类型：

数组