# Path: vmoe/train/optimizer.py
# def add_decayed_weights(
#     weight_decay: Optional[WeightDecay]) -> optax.GradientTransformation:
#     def weight_decay_fn(key):
#     def weight_decay_fn(unused_key):
#   def init_fn(_):
#   def update_fn(updates, state, params):
# def create_optimizer(
#     *,
#     name: str,
#     total_steps: int,
#     learning_rate: Union[float, Mapping[str, Any]],
#     gradient_clip: Optional[Mapping[str, Any]] = None,
#     weight_decay: Optional[WeightDecay] = None,
#     frozen_pattern: Optional[Union[str, Sequence[str]]] = None,
#     trainable_pattern: Optional[Union[str, Sequence[str]]] = None,
#     **optim_kwargs) -> optax.GradientTransformation:
# def freeze_weights(
#     *,
#     frozen_pattern: Optional[Union[str, Sequence[str]]],
#     trainable_pattern: Optional[Union[str, Sequence[str]]],
# ) -> optax.GradientTransformation:
#   def frozen_fn(params: optax.Params) -> PyTree:
# def gradient_clipping(
#     global_norm: Optional[float] = None,
#     absolute_value: Optional[float] = None) -> optax.GradientTransformation:
# def trace_momentum(momentum: Optional[float] = None, **kwargs):


import math
import chex
import jax
import jax.numpy as jnp
from absl.testing import absltest
from absl.testing import parameterized
from vmoe.train import optimizer
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Tests for optimizer."""



class AddDecayedWeightsTest(absltest.TestCase):

  def test_fixed_factor(self):
    init_fn, update_fn = optimizer.add_decayed_weights(0.1)
    state = init_fn('foo')
    self.assertIsInstance(state, optimizer.optax.AddDecayedWeightsState)
    grads = {'a': 0.1, 'b': 0.2}
    params = {'a': 0.1, 'b': 0.2}
    new_grads, new_state = update_fn(grads, state, params)
    self.assertIs(new_state, state)
    chex.assert_trees_all_close(
        new_grads, {'a': 0.1 + 0.1 * 0.1, 'b': 0.2 + 0.1 * 0.2}, rtol=1e-5)

  def test_variable_factor(self):
    weight_decay = [('a', 0.1), ('b', 0.2), ('.*z', 0.3)]
    grads = {'a': 0, 'bz': 0, 'cz': 0, 'd': 1}
    params = {'a': 1, 'bz': 2, 'cz': 3, 'd': 4}

# Path: vmoe/train/optimizer.py
# def add_decayed_weights(
#     weight_decay: Optional[WeightDecay]) -> optax.GradientTransformation:
#     def weight_decay_fn(key):
#     def weight_decay_fn(unused_key):
#   def init_fn(_):
#   def update_fn(updates, state, params):
# def create_optimizer(
#     *,
#     name: str,
#     total_steps: int,
#     learning_rate: Union[float, Mapping[str, Any]],
#     gradient_clip: Optional[Mapping[str, Any]] = None,
#     weight_decay: Optional[WeightDecay] = None,
#     frozen_pattern: Optional[Union[str, Sequence[str]]] = None,
#     trainable_pattern: Optional[Union[str, Sequence[str]]] = None,
#     **optim_kwargs) -> optax.GradientTransformation:
# def freeze_weights(
#     *,
#     frozen_pattern: Optional[Union[str, Sequence[str]]],
#     trainable_pattern: Optional[Union[str, Sequence[str]]],
# ) -> optax.GradientTransformation:
#   def frozen_fn(params: optax.Params) -> PyTree:
# def gradient_clipping(
#     global_norm: Optional[float] = None,
#     absolute_value: Optional[float] = None) -> optax.GradientTransformation:
# def trace_momentum(momentum: Optional[float] = None, **kwargs):


# Path: vmoe/train/optimizer_test.py
import math
import chex
import jax
import jax.numpy as jnp
from absl.testing import absltest
from absl.testing import parameterized
from vmoe.train import optimizer

#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Tests for optimizer."""



class AddDecayedWeightsTest(absltest.TestCase):

  def test_fixed_factor(self):
    init_fn, update_fn = optimizer.add_decayed_weights(0.1)
    state = init_fn('foo')
    self.assertIsInstance(state, optimizer.optax.AddDecayedWeightsState)
    grads = {'a': 0.1, 'b': 0.2}
    params = {'a': 0.1, 'b': 0.2}
    new_grads, new_state = update_fn(grads, state, params)
    self.assertIs(new_state, state)
    chex.assert_trees_all_close(
        new_grads, {'a': 0.1 + 0.1 * 0.1, 'b': 0.2 + 0.1 * 0.2}, rtol=1e-5)

  def test_variable_factor(self):
    weight_decay = [('a', 0.1), ('b', 0.2), ('.*z', 0.3)]
    grads = {'a': 0, 'bz': 0, 'cz': 0, 'd': 1}
    params = {'a': 1, 'bz': 2, 'cz': 3, 'd': 4}

    init_fn, update_fn = optimizer.add_decayed_weights(weight_decay)