# Path: fedjax/core/typing.py



from typing import Iterable, Tuple
from fedjax.core.typing import PyTree
import jax
import jax.numpy as jnp
"""Utilities for working with tree-like container data structures.

In JAX, the term pytree refers to a tree-like structure built out of
container-like Python objects.
For more details, see https://jax.readthedocs.io/en/latest/pytrees.html.
"""





@jax.jit
def tree_weight(pytree: PyTree, weight: float) -> PyTree:
  """Weights tree leaves by weight."""
  return jax.tree_map(lambda l: l * weight, pytree)


def tree_inverse_weight(pytree: PyTree, weight: float) -> PyTree:
  """Weights tree leaves by ``1 / weight``."""
  inverse_weight = (1. / weight) if weight > 0. else 0.
  return tree_weight(pytree, inverse_weight)


@jax.jit
def tree_zeros_like(pytree: PyTree) -> PyTree:
  """Creates a tree with zeros with same structure as the input."""
  return jax.tree_map(jnp.zeros_like, pytree)


@jax.jit

# Path: fedjax/core/typing.py



# Path: fedjax/core/tree_util.py
from typing import Iterable, Tuple
from fedjax.core.typing import PyTree
import jax
import jax.numpy as jnp

"""Utilities for working with tree-like container data structures.

In JAX, the term pytree refers to a tree-like structure built out of
container-like Python objects.
For more details, see https://jax.readthedocs.io/en/latest/pytrees.html.
"""





@jax.jit
def tree_weight(pytree: PyTree, weight: float) -> PyTree:
  """Weights tree leaves by weight."""
  return jax.tree_map(lambda l: l * weight, pytree)


def tree_inverse_weight(pytree: PyTree, weight: float) -> PyTree:
  """Weights tree leaves by ``1 / weight``."""
  inverse_weight = (1. / weight) if weight > 0. else 0.
  return tree_weight(pytree, inverse_weight)


@jax.jit
def tree_zeros_like(pytree: PyTree) -> PyTree:
  """Creates a tree with zeros with same structure as the input."""
  return jax.tree_map(jnp.zeros_like, pytree)


@jax.jit

def tree_add(left: PyTree, right: PyTree) -> PyTree: