# 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
                client_init=client_init,
                shared_input=shared_input,
                client_input=client_input) from e
          for batch in client_batches:
            try:
              state, step_result = client_step(state, batch)
            except Exception as e:
              raise ForEachClientError(
                  e,
                  stage='client_step',
                  client_id=client_id,
                  client_step=client_step,
                  state=state,
                  batch=batch) from e
            step_results.append(step_result)
          try:
            output = client_final(shared_input, state)
          except Exception as e:
            raise ForEachClientError(
                e,
                stage='client_final',
                client_id=client_id,
                client_final=client_final,
                shared_input=shared_input,
                state=state) from e
          yield client_id, output, step_results

    return run



# 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

                client_init=client_init,
                shared_input=shared_input,
                client_input=client_input) from e
          for batch in client_batches:
            try:
              state, step_result = client_step(state, batch)
            except Exception as e:
              raise ForEachClientError(
                  e,
                  stage='client_step',
                  client_id=client_id,
                  client_step=client_step,
                  state=state,
                  batch=batch) from e
            step_results.append(step_result)
          try:
            output = client_final(shared_input, state)
          except Exception as e:
            raise ForEachClientError(
                e,
                stage='client_final',
                client_id=client_id,
                client_final=client_final,
                shared_input=shared_input,
                state=state) from e
          yield client_id, output, step_results

    return run



@dataclasses.dataclass