# Path: fedjax/core/dataclasses.py
# def dataclass(clz: type):
#   def replace(self, **updates):
#   def iterate_clz(x):
#   def clz_from_iterable(meta, data):
# 
# Path: fedjax/core/typing.py

# 
# Path: fedjax/core/typing.py



from typing import Callable, List, Optional, Tuple, Union
from fedjax.core import dataclasses
from fedjax.core.typing import OptState
from fedjax.core.typing import Params
import haiku as hk
import jax
import optax



Grads = Params


@dataclasses.dataclass
class Optimizer:
  """Wraps different optimizer libraries in a common interface.

  Works with `optax <https://github.com/deepmind/optax>`_.

  The expected usage of Optimizer is as follows::

    # One step of SGD.
    params = {'w': jnp.array([1, 1, 1])}
    grads = {'w': jnp.array([2, 3, 4])}
    optimizer = fedjax.optimizers.sgd(learning_rate=0.1)
    opt_state = optimizer.init(params)
    opt_state, params = optimizer.apply(grads, opt_state, params)
    print(params)
    # {'w': DeviceArray([0.8, 0.7, 0.6], dtype=float32)}

  Attributes:
    init: Initializes (possibly empty) PyTree of statistics (optimizer state)
      given the input model parameters.
    apply: Transforms and applies the input gradients to update the optimizer
      state and model parameters.
  """
  init: Callable[[Params], OptState]

# Path: fedjax/core/dataclasses.py
# def dataclass(clz: type):
#   def replace(self, **updates):
#   def iterate_clz(x):
#   def clz_from_iterable(meta, data):
# 
# Path: fedjax/core/typing.py

# 
# Path: fedjax/core/typing.py



# Path: fedjax/core/optimizers.py
from typing import Callable, List, Optional, Tuple, Union
from fedjax.core import dataclasses
from fedjax.core.typing import OptState
from fedjax.core.typing import Params
import haiku as hk
import jax
import optax




Grads = Params


@dataclasses.dataclass
class Optimizer:
  """Wraps different optimizer libraries in a common interface.

  Works with `optax <https://github.com/deepmind/optax>`_.

  The expected usage of Optimizer is as follows::

    # One step of SGD.
    params = {'w': jnp.array([1, 1, 1])}
    grads = {'w': jnp.array([2, 3, 4])}
    optimizer = fedjax.optimizers.sgd(learning_rate=0.1)
    opt_state = optimizer.init(params)
    opt_state, params = optimizer.apply(grads, opt_state, params)
    print(params)
    # {'w': DeviceArray([0.8, 0.7, 0.6], dtype=float32)}

  Attributes:
    init: Initializes (possibly empty) PyTree of statistics (optimizer state)
      given the input model parameters.
    apply: Transforms and applies the input gradients to update the optimizer
      state and model parameters.
  """
  init: Callable[[Params], OptState]

  apply: Callable[[Grads, OptState, Params], Tuple[OptState, Params]]