broadcasted_iota

Contents

broadcasted_iota#

class brainunit.lax.broadcasted_iota(dtype, shape, dimension, _sharding=None, unit=None, **kwargs)#

Broadcast an iota array into the given shape along one dimension.

Convenience wrapper around iota.

Parameters:

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

  • shape (int | Sequence[int]) – The shape of the output array.

  • dimension (int) – The dimension along which to broadcast the iota values.

  • _sharding (optional) – Internal sharding parameter.

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

Returns:

result – An array of the given shape with iota values along dimension.

Return type:

Quantity | Array

Examples

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