# 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
# Licensed under the Apache License, Version 2.0 (the "License");
# 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.
"""Transformations and utilities for specifying work for each client.

Many of these utilities work with and require PyTrees as input and output.
This is required when working with the underlying JAX transformations like
`jax.pmap` and `jax.jit`. For more information on what PyTrees are, refer to
https://jax.readthedocs.io/en/latest/pytrees.html.
"""




# 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.

# 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

# Licensed under the Apache License, Version 2.0 (the "License");
# 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.
"""Transformations and utilities for specifying work for each client.

Many of these utilities work with and require PyTrees as input and output.
This is required when working with the underlying JAX transformations like
`jax.pmap` and `jax.jit`. For more information on what PyTrees are, refer to
https://jax.readthedocs.io/en/latest/pytrees.html.
"""




# 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