jax.lax.broadcasted_iota

jax.lax.broadcasted_iota(dtype, shape, dimension, *, out_sharding=None)

围绕 iota 的便捷封装器。

参数:

• dtype (DTypeLike)

• shape (Shape)

• dimension (int)

返回类型:

Array