flax
Make MultiOptimizer use apply_gradient instead of apply_param_gradient
#1195
Merged
Merged