qr

Contents

qr#

class brainunit.lax.qr(x, *, pivoting=False, full_matrices=True, use_magma=None)#

QR decomposition.

Compute the QR decomposition \(A = Q \cdot R\) where \(Q\) is a unitary (orthogonal) matrix and \(R\) is upper-triangular.

Parameters:

  • x (Union[saiunit.Quantity, Array, ndarray, number, bool]) – A batch of matrices with shape [..., m, n].

  • pivoting (bool) – If True, compute a column-pivoted decomposition \(A[:, P] = Q \cdot R\) and additionally return the pivot indices p. Default is False.

  • full_matrices (bool) – If True, compute the full-size q and r; if False, only the leading min(m, n) columns of q are returned. Default is True.

  • use_magma (bool | None) – Whether to use MAGMA for the pivoted decomposition on GPU. Default is None.

Return type:

tuple[Array, Union[saiunit.Quantity, Array]] | tuple[Array, Union[saiunit.Quantity, Array], Array]

Returns:

  • q (Array) – The unitary factor (unitless).

  • r (Array or Quantity) – The upper-triangular factor with shape [..., min(m, n), n]. If x has a unit, r preserves that unit.

  • p (Array) – Column pivot indices (plain integers). Only returned when pivoting is True.

Examples

>>> import jax.numpy as jnp
>>> import saiunit as u
>>> import saiunit.lax as sulax
>>> A = jnp.array([[1.0, 2.0], [3.0, 4.0]]) * u.meter
>>> q, r = sulax.qr(A)
>>> u.get_unit(r) == u.meter
True
Contents