Efficient Lax Triangular Solve For A Large Matrix

[adam-hartshorne]

I am looking for for an efficient way to calculate jax.lax.linalg.triangular_solve(L, b), where L = [N N] and b is [M N D], given that M is large, thus I don’t want to make L = [M N N] by copying it using a repeat.

Any help much appreciated.

[jakevdp]

If I understand your question correctly, I think you can do what you want with jax.vmap. Here is a short example:

import jax.numpy as jnp
from jax import random, vmap
from jax.scipy.linalg import solve_triangular

N = 10
M = 100
D = 5

key = random.PRNGKey(1701)
key1, key2 = random.split(key)

L = jnp.triu(random.uniform(key1, (N, N)))
b = random.uniform(key2, (M, N, D))

# slow version: call solve_triangular once per entry in b & stack the result
out_slow = jnp.stack([
  solve_triangular(L, b[i]) for i in range(M)
])

# fast version: vmap with appropriate in_axes
out = vmap(solve_triangular, in_axes=(None, 0))(L, b)

print(jnp.allclose(out, out_slow))
# True

[adam-hartshorne]

Thanks, that's perfect.

[hawkinsp]

A thing to note is that triangular solves are usually written already batched on the right-hand side. The algorithm for a matrix RHS is just a batched version of the algorithm with a vector RHS. So you can add another batch dimension by folding it into one of the existing matrix dimensions.

But happily this knowledge is already encoded in the JAX batching rule for triangular solves so you don't need to know that
(https://cs.opensource.google/jax/jax/+/main:jax/_src/lax/linalg.py;l=633?q=jax%20triangular_solve_batching_rule), just use vmap and be happy.