# 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/shakespeare.py
# def create_lstm_model(vocab_size: int = 86,
#                       embed_size: int = 8,
#                       lstm_hidden_size: int = 256,
#                       lstm_num_layers: int = 2) -> 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 shakespeare
import jax
import jax.numpy as jnp
# Copyright 2021 Google LLC
#
# 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.
"""Tests for fedjax.models.shakespeare."""






class ShakespeareModelTest(absltest.TestCase):

  def test_create_lstm_model(self):
    model = shakespeare.create_lstm_model()
    params = model.init(jax.random.PRNGKey(0))
    batch = {
        'x': jnp.ones((5, 3), dtype=jnp.int32),
        'y': jnp.ones((5, 3), dtype=jnp.int32)
    }

# 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/shakespeare.py
# def create_lstm_model(vocab_size: int = 86,
#                       embed_size: int = 8,
#                       lstm_hidden_size: int = 256,
#                       lstm_num_layers: int = 2) -> models.Model:
#   def forward_pass(batch):
#   def train_loss(batch, preds):


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

# Copyright 2021 Google LLC
#
# 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.
"""Tests for fedjax.models.shakespeare."""






class ShakespeareModelTest(absltest.TestCase):

  def test_create_lstm_model(self):
    model = shakespeare.create_lstm_model()
    params = model.init(jax.random.PRNGKey(0))
    batch = {
        'x': jnp.ones((5, 3), dtype=jnp.int32),
        'y': jnp.ones((5, 3), dtype=jnp.int32)
    }

    self.assertEqual(tree_util.tree_size(params), 820522)