# 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/stackoverflow.py
# def create_lstm_model(vocab_size: int = 10000,
#                       embed_size: int = 96,
#                       lstm_hidden_size: int = 670,
#                       lstm_num_layers: int = 1,
#                       share_input_output_embeddings: bool = False,
#                       expected_length: Optional[float] = None) -> models.Model:
#   def forward_pass(batch):
#   def train_loss(batch, preds):


from typing import Hashable
from absl.testing import absltest
from fedjax.core import tree_util
from fedjax.models import stackoverflow
import jax
import jax.numpy as jnp
import numpy.testing as npt
    }
    with self.subTest('apply_for_train'):
      preds = model.apply_for_train(params, batch)
      self.assertTupleEqual(preds.shape, (5, 3, 10000 + 4))
    with self.subTest('apply_for_eval'):
      preds = model.apply_for_eval(params, batch)
      self.assertTupleEqual(preds.shape, (5, 3, 10000 + 4))
    with self.subTest('train_loss'):
      preds = model.apply_for_train(params, batch)
      train_loss = model.train_loss(batch, preds)
      self.assertTupleEqual(train_loss.shape, (5,))
    with self.subTest('hashable'):
      self.assertIsInstance(model, Hashable)

  def test_create_lstm_model(self):
    model = stackoverflow.create_lstm_model()
    params = model.init(jax.random.PRNGKey(0))
    self.assertEqual(tree_util.tree_size(params), 4050748)
    self.check_model(model)

  def test_create_lstm_model_share_embeddings(self):
    model = stackoverflow.create_lstm_model(share_input_output_embeddings=True)
    params = model.init(jax.random.PRNGKey(0))
    self.assertEqual(tree_util.tree_size(params), 3090364)
    self.check_model(model)

  def test_expected_length_scale_loss(self):
    batch = {'y': jnp.ones((5, 3), dtype=jnp.int32)}
    preds = jnp.ones((5, 3, 10000))
    expected_length = 3.

# 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/stackoverflow.py
# def create_lstm_model(vocab_size: int = 10000,
#                       embed_size: int = 96,
#                       lstm_hidden_size: int = 670,
#                       lstm_num_layers: int = 1,
#                       share_input_output_embeddings: bool = False,
#                       expected_length: Optional[float] = None) -> models.Model:
#   def forward_pass(batch):
#   def train_loss(batch, preds):


# Path: fedjax/models/stackoverflow_test.py
from typing import Hashable
from absl.testing import absltest
from fedjax.core import tree_util
from fedjax.models import stackoverflow
import jax
import jax.numpy as jnp
import numpy.testing as npt

    }
    with self.subTest('apply_for_train'):
      preds = model.apply_for_train(params, batch)
      self.assertTupleEqual(preds.shape, (5, 3, 10000 + 4))
    with self.subTest('apply_for_eval'):
      preds = model.apply_for_eval(params, batch)
      self.assertTupleEqual(preds.shape, (5, 3, 10000 + 4))
    with self.subTest('train_loss'):
      preds = model.apply_for_train(params, batch)
      train_loss = model.train_loss(batch, preds)
      self.assertTupleEqual(train_loss.shape, (5,))
    with self.subTest('hashable'):
      self.assertIsInstance(model, Hashable)

  def test_create_lstm_model(self):
    model = stackoverflow.create_lstm_model()
    params = model.init(jax.random.PRNGKey(0))
    self.assertEqual(tree_util.tree_size(params), 4050748)
    self.check_model(model)

  def test_create_lstm_model_share_embeddings(self):
    model = stackoverflow.create_lstm_model(share_input_output_embeddings=True)
    params = model.init(jax.random.PRNGKey(0))
    self.assertEqual(tree_util.tree_size(params), 3090364)
    self.check_model(model)

  def test_expected_length_scale_loss(self):
    batch = {'y': jnp.ones((5, 3), dtype=jnp.int32)}
    preds = jnp.ones((5, 3, 10000))
    expected_length = 3.

    model = stackoverflow.create_lstm_model()