Primitive Triangular Solve

Solves a system of linear equations with a triangular coefficient matrix. When left_side is TRUE, solves op(a) %*% x = b for x. When left_side is FALSE, solves x %*% op(a) = b for x. Axes before the last two are batch axes and must match between a and b (no broadcasting). Here op is A or A^T depending on transpose_a.

Usage

prim_triangular_solve(a, b, left_side, lower, unit_diagonal, transpose_a)

Arguments

a: (arrayish)
Triangular coefficient matrix with at least 2 axes. The last two axes must be equal (square matrix); any leading axes are batch axes. Can be any float data type. a and b must have the same data type. An R value among them assumes the data type of the others when it is in its data type category, and its default data type when none of them has one.
b: (arrayish)
Right-hand side. Same rank as a (rank >= 2), with matching leading batch axes. The size of a's last two (square) axes must equal b's second-to-last axis when left_side = TRUE, or b's last axis when left_side = FALSE. Shares a's data type – see a.
left_side: (logical(1))
If TRUE, solve op(a) %*% x = b. If FALSE, solve x %*% op(a) = b.
lower: (logical(1))
If TRUE, a is lower triangular. If FALSE, a is upper triangular.
unit_diagonal: (logical(1))
If TRUE, assume diagonal elements of a are 1.
transpose_a: (logical(1))
If TRUE, solve with t(a) in place of a.

Value

(arrayish)
Has b's shape and the data type a and b agreed on.

Details

Differentiation is only implemented for a single system: the reverse rule errors on batched operands (operands with more than 2 axes).

Implemented Rules

stablehlo
reverse

StableHLO

Lowers to hlo_triangular_solve(), specified under triangular_solve.

References

Giles M (2008). “An extended collection of matrix derivative results for forward and reverse mode automatic differentiation.” Oxford University Computing Laboratory.

Examples

# solve L %*% x = b where L is lower triangular
L <- nv_matrix(c(2, 1, 0, 3), nrow = 2, dtype = "f32")
b <- nv_matrix(c(4, 3), nrow = 2, dtype = "f32")
prim_triangular_solve(L, b,
  left_side = TRUE, lower = TRUE,
  unit_diagonal = FALSE, transpose_a = FALSE
)
#> AnvlArray
#>  2.0000
#>  0.3333
#> [ CPUf32{2,1} ]