nifty.re.tree_math.util module#

logm(M)[source]#

Computes the matrix logarithm of a symmetric square matrix. Input matrix must have strictly positive eigenvalues. Mapped over the leafs of the input tree and batched over leading axes.

Parameters:

M (tree-like structure of jnp.ndarray) – The input matrix/matrices

Returns:

out – Matrix logarithm of M

Return type:

tree-like structure of jnp.ndarray (like M)

solve(A, B, *, matrix_eqn=False, transposed=False)[source]#

Solves the linear system of equations

\[AX = B\]

for X, where A is a symmetric positive definite matrix. Mapped over the leafs of the input trees and batched over leading axes. This implementation features a custom, more stable jvp rule.

Parameters:

  • A (tree-like structure of jnp.ndarray) – LHS of the equation

  • B (tree-like structure of jnp.ndarray) – RHS of the equation

  • matrix_eqn (bool) – Specifies whether X and B are matrices or Vectors (default: False)

  • transposed (bool) –

    Solves the transposed system of equations

    \[A^T X^T = B^T\]

    instead. (default: False)

Returns:

out – Solution X of the above equation

Return type:

tree-like structure of jnp.ndarray (like B)

sqrtm(M)[source]#

Computes the principal matrix squre root B of a symmetric square matrix M:

\[M = BB\]

with B being positive semidefinite and symmetric. Mapped over the leafs of the input tree and batched over leading axes. This implementation features a custom, more stable jvp rule (only 1st derivative).

Parameters:

M (tree-like structure of jnp.ndarray) – The input matrix/matrices

Returns:

B – Principal matrix square root of M

Return type:

tree-like structure of jnp.ndarray (like M)