flax

0366d3fc - Make MultiOptimizer use apply_gradient instead of apply_param_gradient

Make MultiOptimizer use apply_gradient instead of apply_param_gradient Some advanced optimizers (like Shampoo) don't define a pure leaf-based optimization scheme in terms of apply_param_gradient for performance reasons. Ensure the new MultiOptimizer works with them. PiperOrigin-RevId: 366957513