flax

d812e604 - don't allow Variable mutation in custom_vjp on differentiable arguments