flax

don't allow Variable mutation in custom_vjp on differentiable arguments

#5342

Merged