def unrolled_substitution_solve(matvec, b, lower_tri):
    'Solve a triangular unrolled system with fwd/back substitution.'
    zero = jnp.zeros(())
    one = jnp.ones(())
    x = [zero for _ in b]
    ordering = (range(len(b)) if lower_tri else range((len(b) - 1), (- 1), (- 1)))
    for i in ordering:
        residual = (b[i] - matvec(x)[i])
        diagonal = matvec([(one if (i == j) else zero) for j in range(len(b))])[i]
        x[i] = (residual / diagonal)
    return x
Solve a triangular unrolled system with fwd/back substitution.
def unrolled_substitution_solve(matvec, b, lower_tri):
    
    zero = jnp.zeros(())
    one = jnp.ones(())
    x = [zero for _ in b]
    ordering = (range(len(b)) if lower_tri else range((len(b) - 1), (- 1), (- 1)))
    for i in ordering:
        residual = (b[i] - matvec(x)[i])
        diagonal = matvec([(one if (i == j) else zero) for j in range(len(b))])[i]
        x[i] = (residual / diagonal)
    return x