flax

check aliases on all transform args and simplify apply_variable_updates

#5349

Merged