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 (Union[str, type[Any], dtype, SupportsDType]) – The element type of the output array.

  • size (int) – The number of elements.

  • unit (Optional[Unit]) – 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:

Union[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