jax.lax.iota

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

Wraps XLA’s Iota operator.

Parameters
  • dtype (Any) –

  • size (int) –

Return type

Any