# 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



import abc
import contextlib
import functools
import threading
import jax
import jax.numpy as jnp
from typing import Any, Callable, Iterable, Iterator, List, Optional, Sequence, Tuple, Union
from fedjax.core import dataclasses
from fedjax.core.typing import BatchExample
from fedjax.core.typing import PyTree
"""




# Shared input that is passed to the client init that is shared across all
# clients. For example, this could be the shared global model parameters that
# are then used as the starting point for per client training.
SharedInput = PyTree
# Client specific input passed into client init that is different for each
# client. For example, this could be the starting PRNG seed for training a given
# client. Additionally, for stateful client federated algorithms (usually in the
# cross-silo setting), this is how per client information will be passed along.
ClientInput = PyTree

# Intermittent client state passed as input and output for each client step.
# This will typically contain model parameters and optimizer state that are
# updated at each client step.
ClientStepState = PyTree
# Step results can be used to record any metrics over the course of running the
# for_each_client loop. For example, it can be used to record train metrics for
# each client like gradient norm.
ClientStepResult = PyTree

# Final output from the for_each_client loop. This is usually a subset of the
# client step state. However, more meaningful transformations can be done like
# model update clipping.
ClientOutput = PyTree

ClientInit = Callable[[SharedInput, ClientInput], ClientStepState]

# 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/for_each_client.py
import abc
import contextlib
import functools
import threading
import jax
import jax.numpy as jnp
from typing import Any, Callable, Iterable, Iterator, List, Optional, Sequence, Tuple, Union
from fedjax.core import dataclasses
from fedjax.core.typing import BatchExample
from fedjax.core.typing import PyTree

"""




# Shared input that is passed to the client init that is shared across all
# clients. For example, this could be the shared global model parameters that
# are then used as the starting point for per client training.
SharedInput = PyTree
# Client specific input passed into client init that is different for each
# client. For example, this could be the starting PRNG seed for training a given
# client. Additionally, for stateful client federated algorithms (usually in the
# cross-silo setting), this is how per client information will be passed along.
ClientInput = PyTree

# Intermittent client state passed as input and output for each client step.
# This will typically contain model parameters and optimizer state that are
# updated at each client step.
ClientStepState = PyTree
# Step results can be used to record any metrics over the course of running the
# for_each_client loop. For example, it can be used to record train metrics for
# each client like gradient norm.
ClientStepResult = PyTree

# Final output from the for_each_client loop. This is usually a subset of the
# client step state. However, more meaningful transformations can be done like
# model update clipping.
ClientOutput = PyTree

ClientInit = Callable[[SharedInput, ClientInput], ClientStepState]

ClientStep = Callable[[ClientStepState, BatchExample], Tuple[ClientStepState,