jax.lax.iota#

jax.lax.iota(dtype, size)[source]#

Wraps XLA’s Iota operator.

Parameters:

  • dtype (DTypeLike)

  • size (int)

Return type:

Array