def test_custom_linear_solve_errors(self):
    solve = (lambda f, x: x)
    with self.assertRaisesRegex(TypeError, re.escape('matvec() output pytree')):
        lax.custom_linear_solve((lambda x: [x]), 1.0, solve, solve)
    with self.assertRaisesRegex(TypeError, re.escape('solve() output pytree')):
        lax.custom_linear_solve((lambda x: x), 1.0, (lambda f, x: [x]), solve)
    with self.assertRaisesRegex(TypeError, re.escape('transpose_solve() output pytree')):
        lax.custom_linear_solve((lambda x: x), 1.0, solve, (lambda f, x: [x]))
    with self.assertRaisesRegex(ValueError, re.escape('solve() output shapes')):
        lax.custom_linear_solve((lambda x: x), 1.0, (lambda f, x: jnp.ones(2)), solve)

    def bad_matvec_usage(a):
        return lax.custom_linear_solve((lambda x: (a * jnp.ones(2))), 1.0, solve, solve)
    with self.assertRaisesRegex(ValueError, re.escape('matvec() output shapes')):
        api.jvp(bad_matvec_usage, (1.0,), (1.0,))
def test_custom_linear_solve_errors(self):
    solve = (lambda f, x: x)
    with self.assertRaisesRegex(TypeError, re.escape('matvec() output pytree')):
        lax.custom_linear_solve((lambda x: [x]), 1.0, solve, solve)
    with self.assertRaisesRegex(TypeError, re.escape('solve() output pytree')):
        lax.custom_linear_solve((lambda x: x), 1.0, (lambda f, x: [x]), solve)
    with self.assertRaisesRegex(TypeError, re.escape('transpose_solve() output pytree')):
        lax.custom_linear_solve((lambda x: x), 1.0, solve, (lambda f, x: [x]))
    with self.assertRaisesRegex(ValueError, re.escape('solve() output shapes')):
        lax.custom_linear_solve((lambda x: x), 1.0, (lambda f, x: jnp.ones(2)), solve)

    def bad_matvec_usage(a):
        return lax.custom_linear_solve((lambda x: (a * jnp.ones(2))), 1.0, solve, solve)
    with self.assertRaisesRegex(ValueError, re.escape('matvec() output shapes')):
        api.jvp(bad_matvec_usage, (1.0,), (1.0,))