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


from absl.testing import absltest
from fedjax.core import tree_util
import jax.numpy as jnp
import numpy.testing as npt
    npt.assert_array_equal(weight_pytree['x'], [[[8.0, 10.0]], [[2.0, 2.0]]])
    npt.assert_array_equal(weight_pytree['y'], [[6.0], [2.0]])

  def test_tree_sum(self):
    pytree_1 = {
        'x': jnp.array([[[4, 5]], [[1, 1]]]),
        'y': jnp.array([[3], [1]]),
    }
    pytree_2 = {
        'x': jnp.array([[[2, 3]], [[4, 5]]]),
        'y': jnp.array([[6], [7]]),
    }
    pytree = tree_util.tree_sum([pytree_1, pytree_2])
    npt.assert_array_equal(pytree['x'], [[[6, 8]], [[5, 6]]])
    npt.assert_array_equal(pytree['y'], [[9], [8]])

  def test_tree_mean(self):
    pytrees = [(0, 1), (2, 3), (4, 5)]
    weights = [6., 7., 8.]
    pytrees_and_weights = zip(pytrees, weights)
    pytree = tree_util.tree_mean(pytrees_and_weights)
    npt.assert_array_almost_equal(pytree,
                                  (2.1904761904761907, 3.1904761904761907))

  def test_tree_clip_by_global_norm(self):
    pytree = {
        'x': jnp.array([[[4, 5]], [[1, 1]]]),
        'y': jnp.array([[3], [1]]),
    }
    max_norm = 3.640055  # 0.5 * tree_l2_norm(pytree)

# 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/core/tree_util_test.py
from absl.testing import absltest
from fedjax.core import tree_util
import jax.numpy as jnp
import numpy.testing as npt

    npt.assert_array_equal(weight_pytree['x'], [[[8.0, 10.0]], [[2.0, 2.0]]])
    npt.assert_array_equal(weight_pytree['y'], [[6.0], [2.0]])

  def test_tree_sum(self):
    pytree_1 = {
        'x': jnp.array([[[4, 5]], [[1, 1]]]),
        'y': jnp.array([[3], [1]]),
    }
    pytree_2 = {
        'x': jnp.array([[[2, 3]], [[4, 5]]]),
        'y': jnp.array([[6], [7]]),
    }
    pytree = tree_util.tree_sum([pytree_1, pytree_2])
    npt.assert_array_equal(pytree['x'], [[[6, 8]], [[5, 6]]])
    npt.assert_array_equal(pytree['y'], [[9], [8]])

  def test_tree_mean(self):
    pytrees = [(0, 1), (2, 3), (4, 5)]
    weights = [6., 7., 8.]
    pytrees_and_weights = zip(pytrees, weights)
    pytree = tree_util.tree_mean(pytrees_and_weights)
    npt.assert_array_almost_equal(pytree,
                                  (2.1904761904761907, 3.1904761904761907))

  def test_tree_clip_by_global_norm(self):
    pytree = {
        'x': jnp.array([[[4, 5]], [[1, 1]]]),
        'y': jnp.array([[3], [1]]),
    }
    max_norm = 3.640055  # 0.5 * tree_l2_norm(pytree)

    clipped_pytree = tree_util.tree_clip_by_global_norm(pytree, max_norm)