# Path: fedjax/core/tree_util.py
# def tree_weight(pytree: PyTree, weight: float) -> PyTree:
# def tree_inverse_weight(pytree: PyTree, weight: float) -> PyTree:
# def tree_zeros_like(pytree: PyTree) -> PyTree:
# def tree_add(left: PyTree, right: PyTree) -> PyTree:
# def tree_sum(pytrees: Iterable[PyTree]) -> PyTree:
# def tree_mean(pytrees_and_weights: Iterable[Tuple[PyTree, float]]) -> PyTree:
# def tree_size(pytree: PyTree) -> int:
# def tree_l2_squared(pytree: PyTree) -> float:
# def tree_l2_norm(pytree: PyTree) -> float:
# def tree_clip_by_global_norm(pytree: PyTree, max_norm: float) -> PyTree:
# 
# Path: fedjax/models/emnist.py
# class Dropout(hk.Module):
# class ConvDropoutModule(hk.Module):
#   def __init__(self, rate: float = 0.5):
#   def __call__(self, x: jnp.ndarray, is_train: bool):
#   def __init__(self, num_classes):
#   def __call__(self, x: jnp.ndarray, is_train: bool):
# def create_conv_model(only_digits: bool = False) -> models.Model:
#   def forward_pass(batch, is_train=True):
# def create_dense_model(only_digits: bool = False,
#                        hidden_units: int = 200) -> models.Model:
#   def forward_pass(batch):
# def create_logistic_model(only_digits: bool = False) -> models.Model:
#   def forward_pass(batch):
# def create_stax_dense_model(only_digits: bool = False,
#                             hidden_units: int = 200) -> models.Model:
# _HAIKU_SAMPLE_BATCH = {
#     'x': np.zeros((1, 28, 28, 1), dtype=np.float32),
#     'y': np.zeros(1, dtype=np.float32)
# }
# _STAX_SAMPLE_SHAPE = (-1, 28, 28, 1)
# _TRAIN_LOSS = lambda b, p: metrics.unreduced_cross_entropy_loss(b['y'], p)
# _EVAL_METRICS = {
#     'loss': metrics.CrossEntropyLoss(),
#     'accuracy': metrics.Accuracy()
# }


from absl.testing import absltest
from fedjax.core import tree_util
from fedjax.models import emnist
import jax
import jax.numpy as jnp
    with self.subTest('apply_for_train'):
      preds = model.apply_for_train(params, batch, rng)
      self.assertTupleEqual(preds.shape, (5, 62))
    with self.subTest('apply_for_eval'):
      preds = model.apply_for_eval(params, batch)
      self.assertTupleEqual(preds.shape, (5, 62))
    with self.subTest('train_loss'):
      preds = model.apply_for_train(params, batch, rng)
      train_loss = model.train_loss(batch, preds)
      self.assertTupleEqual(train_loss.shape, (5,))

  def test_create_conv_model(self):
    model = emnist.create_conv_model(only_digits=False)
    params = model.init(jax.random.PRNGKey(0))
    self.assertEqual(tree_util.tree_size(params), 1206590)
    self.check_model(model)

  def test_create_dense_model(self):
    model = emnist.create_dense_model(only_digits=False, hidden_units=200)
    params = model.init(jax.random.PRNGKey(0))
    self.assertEqual(tree_util.tree_size(params), 209662)
    self.check_model(model)

  def test_create_logistic_model(self):
    model = emnist.create_logistic_model(only_digits=False)
    params = model.init(jax.random.PRNGKey(0))
    self.assertEqual(tree_util.tree_size(params), 48670)
    self.check_model(model)

  def test_create_stax_dense_model(self):

# Path: fedjax/core/tree_util.py
# def tree_weight(pytree: PyTree, weight: float) -> PyTree:
# def tree_inverse_weight(pytree: PyTree, weight: float) -> PyTree:
# def tree_zeros_like(pytree: PyTree) -> PyTree:
# def tree_add(left: PyTree, right: PyTree) -> PyTree:
# def tree_sum(pytrees: Iterable[PyTree]) -> PyTree:
# def tree_mean(pytrees_and_weights: Iterable[Tuple[PyTree, float]]) -> PyTree:
# def tree_size(pytree: PyTree) -> int:
# def tree_l2_squared(pytree: PyTree) -> float:
# def tree_l2_norm(pytree: PyTree) -> float:
# def tree_clip_by_global_norm(pytree: PyTree, max_norm: float) -> PyTree:
# 
# Path: fedjax/models/emnist.py
# class Dropout(hk.Module):
# class ConvDropoutModule(hk.Module):
#   def __init__(self, rate: float = 0.5):
#   def __call__(self, x: jnp.ndarray, is_train: bool):
#   def __init__(self, num_classes):
#   def __call__(self, x: jnp.ndarray, is_train: bool):
# def create_conv_model(only_digits: bool = False) -> models.Model:
#   def forward_pass(batch, is_train=True):
# def create_dense_model(only_digits: bool = False,
#                        hidden_units: int = 200) -> models.Model:
#   def forward_pass(batch):
# def create_logistic_model(only_digits: bool = False) -> models.Model:
#   def forward_pass(batch):
# def create_stax_dense_model(only_digits: bool = False,
#                             hidden_units: int = 200) -> models.Model:
# _HAIKU_SAMPLE_BATCH = {
#     'x': np.zeros((1, 28, 28, 1), dtype=np.float32),
#     'y': np.zeros(1, dtype=np.float32)
# }
# _STAX_SAMPLE_SHAPE = (-1, 28, 28, 1)
# _TRAIN_LOSS = lambda b, p: metrics.unreduced_cross_entropy_loss(b['y'], p)
# _EVAL_METRICS = {
#     'loss': metrics.CrossEntropyLoss(),
#     'accuracy': metrics.Accuracy()
# }


# Path: fedjax/models/emnist_test.py
from absl.testing import absltest
from fedjax.core import tree_util
from fedjax.models import emnist
import jax
import jax.numpy as jnp

    with self.subTest('apply_for_train'):
      preds = model.apply_for_train(params, batch, rng)
      self.assertTupleEqual(preds.shape, (5, 62))
    with self.subTest('apply_for_eval'):
      preds = model.apply_for_eval(params, batch)
      self.assertTupleEqual(preds.shape, (5, 62))
    with self.subTest('train_loss'):
      preds = model.apply_for_train(params, batch, rng)
      train_loss = model.train_loss(batch, preds)
      self.assertTupleEqual(train_loss.shape, (5,))

  def test_create_conv_model(self):
    model = emnist.create_conv_model(only_digits=False)
    params = model.init(jax.random.PRNGKey(0))
    self.assertEqual(tree_util.tree_size(params), 1206590)
    self.check_model(model)

  def test_create_dense_model(self):
    model = emnist.create_dense_model(only_digits=False, hidden_units=200)
    params = model.init(jax.random.PRNGKey(0))
    self.assertEqual(tree_util.tree_size(params), 209662)
    self.check_model(model)

  def test_create_logistic_model(self):
    model = emnist.create_logistic_model(only_digits=False)
    params = model.init(jax.random.PRNGKey(0))
    self.assertEqual(tree_util.tree_size(params), 48670)
    self.check_model(model)

  def test_create_stax_dense_model(self):

    model = emnist.create_stax_dense_model(only_digits=False, hidden_units=200)