jit grad update and donate args properly

[fehiepsi]

This is useful for optimizing large parameters, like in LLM.

Also update for the new jax.tree pattern.