# Path: fedjax/core/optimizers.py
# class Optimizer:
# def create_optimizer_from_optax(opt: optax.GradientTransformation) -> Optimizer:
#   def apply(grads, opt_state, params):
# def ignore_grads_haiku(optimizer: Optimizer,
#                        non_trainable_names: List[Tuple[str, str]]) -> Optimizer:
#   def non_trainable_to_none(module_name, name, value):
#   def init(params):
#   def apply(grads, opt_state, params):
# def adagrad(learning_rate: ScalarOrSchedule,
#             initial_accumulator_value: float = 0.1,
#             eps: float = 1e-6) -> Optimizer:
# def adam(learning_rate: ScalarOrSchedule,
#          b1: float = 0.9,
#          b2: float = 0.999,
#          eps: float = 1e-8,
#          eps_root: float = 0.0) -> Optimizer:
# def rmsprop(learning_rate: ScalarOrSchedule,
#             decay: float = 0.9,
#             eps: float = 1e-8,
#             initial_scale: float = 0.,
#             centered: bool = False,
#             momentum: Optional[float] = None,
#             nesterov: bool = False) -> Optimizer:
# def sgd(learning_rate: ScalarOrSchedule,
#         momentum: Optional[float] = None,
#         nesterov: bool = False) -> Optimizer:


from absl.testing import absltest
from fedjax.core import optimizers
import flax
import haiku as hk
import jax
import jax.numpy as jnp
import numpy.testing as npt
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      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 fedjax.core.optimizers."""





class OptimizersTest(absltest.TestCase):

  def test_ignore_grads_haiku(self):
    params = hk.data_structures.to_immutable_dict({
        'linear_1': {
            'w': jnp.array([1., 1., 1.])
        },
        'linear_2': {
            'w': jnp.array([2., 2., 2.]),
            'b': jnp.array([3., 3., 3.])
        }
    })
    grads = jax.tree_util.tree_map(lambda _: 0.5, params)
    ignore_optimizer = optimizers.ignore_grads_haiku(

# Path: fedjax/core/optimizers.py
# class Optimizer:
# def create_optimizer_from_optax(opt: optax.GradientTransformation) -> Optimizer:
#   def apply(grads, opt_state, params):
# def ignore_grads_haiku(optimizer: Optimizer,
#                        non_trainable_names: List[Tuple[str, str]]) -> Optimizer:
#   def non_trainable_to_none(module_name, name, value):
#   def init(params):
#   def apply(grads, opt_state, params):
# def adagrad(learning_rate: ScalarOrSchedule,
#             initial_accumulator_value: float = 0.1,
#             eps: float = 1e-6) -> Optimizer:
# def adam(learning_rate: ScalarOrSchedule,
#          b1: float = 0.9,
#          b2: float = 0.999,
#          eps: float = 1e-8,
#          eps_root: float = 0.0) -> Optimizer:
# def rmsprop(learning_rate: ScalarOrSchedule,
#             decay: float = 0.9,
#             eps: float = 1e-8,
#             initial_scale: float = 0.,
#             centered: bool = False,
#             momentum: Optional[float] = None,
#             nesterov: bool = False) -> Optimizer:
# def sgd(learning_rate: ScalarOrSchedule,
#         momentum: Optional[float] = None,
#         nesterov: bool = False) -> Optimizer:


# Path: fedjax/core/optimizers_test.py
from absl.testing import absltest
from fedjax.core import optimizers
import flax
import haiku as hk
import jax
import jax.numpy as jnp
import numpy.testing as npt

# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      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 fedjax.core.optimizers."""





class OptimizersTest(absltest.TestCase):

  def test_ignore_grads_haiku(self):
    params = hk.data_structures.to_immutable_dict({
        'linear_1': {
            'w': jnp.array([1., 1., 1.])
        },
        'linear_2': {
            'w': jnp.array([2., 2., 2.]),
            'b': jnp.array([3., 3., 3.])
        }
    })
    grads = jax.tree_util.tree_map(lambda _: 0.5, params)
    ignore_optimizer = optimizers.ignore_grads_haiku(

        optimizer=optimizers.sgd(learning_rate=1.0),