def triangular_solve_jvp_rule_a(g_a, ans, a, b, left_side, lower, transpose_a,
conjugate_a, unit_diagonal):
m, n = b.shape[-2:]
k = 1 if unit_diagonal else 0
g_a = np.tril(g_a, k=-k) if lower else np.triu(g_a, k=k)
g_a = lax.neg(g_a)
g_a = np.swapaxes(g_a, -1, -2) if transpose_a else g_a
g_a = np.conj(g_a) if conjugate_a else g_a
dot = partial(lax.dot if g_a.ndim == 2 else lax.batch_matmul,
precision=lax.Precision.HIGHEST)
def a_inverse(rhs):
return triangular_solve(a, rhs, left_side, lower, transpose_a,
conjugate_a, unit_diagonal)
# triangular_solve is about the same cost as matrix multplication (~n^2 FLOPs
# for matrix/vector inputs). Order these operations in whichever order is
# cheaper.
if left_side:
assert g_a.shape[-2:] == a.shape[-2:] == (m, m) and ans.shape[-2:] == (
m, n)
if m > n:
return a_inverse(dot(g_a, ans)) # A^{-1} (∂A X)
else:
return dot(a_inverse(g_a), ans) # (A^{-1} ∂A) X
else:
assert g_a.shape[-2:] == a.shape[-2:] == (n, n) and ans.shape[-2:] == (
m, n)
if m < n:
return a_inverse(dot(ans, g_a)) # (X ∂A) A^{-1}
else:
return dot(ans, a_inverse(g_a)) # X (∂A A^{-1})
def triangular_solve_jvp_rule_a(g_a, ans, a, b, left_side, lower, transpose_a,
conjugate_a):
g_a = lax.neg(g_a)
g_a = np.swapaxes(g_a, -1, -2) if transpose_a else g_a
tmp = triangular_solve(a, g_a, left_side, lower, transpose_a, conjugate_a)
dot = lax.dot if g_a.ndim == 2 else lax.batch_matmul
if left_side:
return dot(tmp, ans)
else:
return dot(ans, tmp)
def triangular_solve_jvp_rule_a(g_a, ans, a, b, left_side, lower, transpose_a,
conjugate_a, unit_diagonal):
k = 1 if unit_diagonal else 0
g_a = np.tril(g_a, k=-k) if lower else np.triu(g_a, k=k)
g_a = lax.neg(g_a)
g_a = np.swapaxes(g_a, -1, -2) if transpose_a else g_a
g_a = np.conj(g_a) if conjugate_a else g_a
tmp = triangular_solve(a, g_a, left_side, lower, transpose_a, conjugate_a,
unit_diagonal)
dot = lax.dot if g_a.ndim == 2 else lax.batch_matmul
if left_side:
return dot(tmp, ans)
else:
return dot(ans, tmp)
def _T(x):
return np.swapaxes(x, -1, -2)
def _T(x): return np.swapaxes(x, -1, -2) def _H(x): return np.conj(_T(x))
