broadcast

Contents

broadcast#

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

Broadcast an array by adding new leading dimensions.

Parameters:
  • operand (saiunit.Quantity | Array | ndarray | bool | number | bool | int | float | complex) – 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:

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