def bad_matvec_usage(a):
    return lax.custom_linear_solve((lambda x: (a * jnp.ones(2))), 1.0, solve, solve)
def bad_matvec_usage(a):
    return lax.custom_linear_solve((lambda x: (a * jnp.ones(2))), 1.0, solve, solve)