jax.lax.iota

Warning

This page was created from a pull request (#9655).

jax.lax.iota¶

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

Wraps XLA’s Iota operator.

Parameters

  • dtype (Any) –

  • size (int) –

Return type

Any

previous

jax.lax.index_take

next

jax.lax.is_finite