broadcast_to_rank

broadcast_to_rank#

class brainunit.lax.broadcast_to_rank(x, rank)[source]#

Add leading dimensions of size 1 to give x rank rank.

Parameters:

  • x (saiunit.Quantity | Array | ndarray | bool | number | bool | int | float | complex) – The input array.

  • rank (int) – The desired rank of the output.

Returns:

result – The array with added leading dimensions. Preserves the unit of x.

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.meter
>>> result = sulax.broadcast_to_rank(q, rank=3)
>>> result.mantissa.shape
(1, 1, 2)