def custom_unrolled_lower_tri_solve(mat, b):
    return lax.custom_linear_solve(partial(unrolled_matvec, mat), b, partial(unrolled_substitution_solve, lower_tri=True), partial(unrolled_substitution_solve, lower_tri=False))
def custom_unrolled_lower_tri_solve(mat, b):
    return lax.custom_linear_solve(partial(unrolled_matvec, mat), b, partial(unrolled_substitution_solve, lower_tri=True), partial(unrolled_substitution_solve, lower_tri=False))