iota

Contents

iota#

class saiunit.lax.iota(dtype, size, unit=None, **kwargs)#

Create an iota array (integer sequence) with an optional unit.

Wraps XLA’s Iota operator.

Parameters:

  • dtype (str | type[Any] | dtype | SupportsDType) – The element type of the output array.

  • size (int) – The number of elements.

  • unit (Unit | None) – If provided, the result is a Quantity with this unit.

Returns:

result – An array [0, 1, 2, ..., size - 1] of the given dtype.

Return type:

Quantity | Array

Examples

>>> import saiunit as u
>>> import saiunit.lax as sulax
>>> result = sulax.iota(float, 5, unit=u.second)
>>> result.mantissa
Array([0., 1., 2., 3., 4.], dtype=float32)
>>> result.unit
second
Contents