broadcast

Contents

broadcast#

class saiunit.lax.broadcast(operand, sizes)[source]#

Broadcast an array by adding new leading dimensions.

Parameters:

  • operand (Union[saiunit.Quantity, Array, ndarray, number, bool]) – The input array.

  • sizes (Sequence[int]) – Sizes of new leading dimensions to prepend to the array shape.

Returns:

result – The broadcasted array. Preserves the unit of operand.

Return type:

Union[saiunit.Quantity, Array]

Examples

>>> import saiunit as u
>>> import saiunit.lax as sulax
>>> import jax.numpy as jnp
>>> q = jnp.array([1.0, 2.0]) * u.second
>>> result = sulax.broadcast(q, sizes=(3,))
>>> result.mantissa.shape
(3, 2)
>>> result.unit
second
Contents