# Path: fedjax/aggregators/compression.py
# class CompressionState:
# def binary_stochastic_quantize(v: jnp.ndarray,
#                                rng: PRNGKey,
#                                v_min: Optional[float] = None,
#                                v_max: Optional[float] = None) -> jnp.ndarray:
# def uniform_stochastic_quantize(v: jnp.ndarray,
#                                 num_levels: int,
#                                 rng: PRNGKey,
#                                 v_min: Optional[float] = None,
#                                 v_max: Optional[float] = None) -> jnp.ndarray:
# def uniform_stochastic_quantize_pytree(params: Params, num_levels: int,
#                                        rng: PRNGKey) -> Params:
# def num_leaves(pytree):
# def _entropy(v, uniq):
# def _hist_bits(v, uniq):
# def arithmetic_encoding_num_bits(v: jnp.ndarray) -> int:
# def uniform_stochastic_quantizer(
#     num_levels: int,
#     rng: PRNGKey,
#     encode_algorithm: Optional[str] = None) -> aggregator.Aggregator:
#   def init():
#   def apply(
#       clients_params_and_weights: Iterable[Tuple[ClientId, Params, float]],
#       aggregator_state: CompressionState) -> Tuple[Params, CompressionState]:
#     def quantize_params_and_weight(client_params_and_weight, rng):
#       def arithmetic_encoding_num_bits_pytree(params, weights):
# def rotated_uniform_stochastic_quantizer(num_levels: int,
#                                          rng: PRNGKey) -> aggregator.Aggregator:
#   def init():
#   def apply(
#       clients_params_and_weights: Iterable[Tuple[ClientId, Params, float]],
#       aggregator_state: CompressionState) -> Tuple[Params, CompressionState]:
#     def quantize_params_and_weight(client_params_and_weight, rng):
# def drive_pytree(params: Params) -> Params:
# def structured_drive_quantizer(rng: PRNGKey) -> aggregator.Aggregator:
#   def init():
#   def apply(
#       clients_params_and_weights: Iterable[Tuple[ClientId, Params, float]],
#       aggregator_state: CompressionState) -> Tuple[Params, CompressionState]:
#     def quantize_params_and_weight(client_id, params, weight):
# def terngrad_quantize(v: jnp.ndarray, rng: PRNGKey) -> jnp.ndarray:
# def terngrad_quantize_pytree(params: Params, rng: PRNGKey) -> Params:
# def terngrad_quantizer(rng: PRNGKey) -> aggregator.Aggregator:
#   def init():
#   def apply(
#       clients_params_and_weights: Iterable[Tuple[ClientId, Params, float]],
#       aggregator_state: CompressionState) -> Tuple[Params, CompressionState]:
#     def quantize_params_and_weight(client_params_and_weight, rng):


from absl.testing import absltest
from fedjax.aggregators import compression
import jax
import jax.numpy as jnp
import numpy.testing as npt
    rng = jax.random.PRNGKey(42)
    compressed_sum = jnp.zeros_like(v)
    for _ in range(500):
      rng, use_rng = jax.random.split(rng)
      compressed_sum += compression.uniform_stochastic_quantize(v, 125, use_rng)
    npt.assert_array_almost_equal(compressed_sum / 500, v, decimal=2)

  def test_uniform_stochastic_quantize_unbiasedness_two_dim(self):
    v = jnp.array([[0., 1., 100.], [0.3, 2.3, 45.]])
    rng = jax.random.PRNGKey(42)
    compressed_sum = jnp.zeros_like(v)
    for _ in range(500):
      rng, use_rng = jax.random.split(rng)
      compressed_sum += compression.uniform_stochastic_quantize(v, 125, use_rng)
    npt.assert_array_almost_equal(compressed_sum / 500, v, decimal=1)

  def test_arithmetic_encoding_num_bits(self):
    v = jnp.array([1.0, 2.0, 3.0, 4.0, 5.0])
    npt.assert_array_almost_equal(
        compression.arithmetic_encoding_num_bits(v), [89.82311], decimal=3)

  def test_uniform_stochastic_quantizer(self):
    delta_params_and_weights = [('a', {
        'w': jnp.array([1., 2., 3.])
    }, 2.), ('b', {
        'w': jnp.array([2., 4., 6.])
    }, 4.), ('c', {
        'w': jnp.array([1., 3., 5.])
    }, 2.)]


# Path: fedjax/aggregators/compression.py
# class CompressionState:
# def binary_stochastic_quantize(v: jnp.ndarray,
#                                rng: PRNGKey,
#                                v_min: Optional[float] = None,
#                                v_max: Optional[float] = None) -> jnp.ndarray:
# def uniform_stochastic_quantize(v: jnp.ndarray,
#                                 num_levels: int,
#                                 rng: PRNGKey,
#                                 v_min: Optional[float] = None,
#                                 v_max: Optional[float] = None) -> jnp.ndarray:
# def uniform_stochastic_quantize_pytree(params: Params, num_levels: int,
#                                        rng: PRNGKey) -> Params:
# def num_leaves(pytree):
# def _entropy(v, uniq):
# def _hist_bits(v, uniq):
# def arithmetic_encoding_num_bits(v: jnp.ndarray) -> int:
# def uniform_stochastic_quantizer(
#     num_levels: int,
#     rng: PRNGKey,
#     encode_algorithm: Optional[str] = None) -> aggregator.Aggregator:
#   def init():
#   def apply(
#       clients_params_and_weights: Iterable[Tuple[ClientId, Params, float]],
#       aggregator_state: CompressionState) -> Tuple[Params, CompressionState]:
#     def quantize_params_and_weight(client_params_and_weight, rng):
#       def arithmetic_encoding_num_bits_pytree(params, weights):
# def rotated_uniform_stochastic_quantizer(num_levels: int,
#                                          rng: PRNGKey) -> aggregator.Aggregator:
#   def init():
#   def apply(
#       clients_params_and_weights: Iterable[Tuple[ClientId, Params, float]],
#       aggregator_state: CompressionState) -> Tuple[Params, CompressionState]:
#     def quantize_params_and_weight(client_params_and_weight, rng):
# def drive_pytree(params: Params) -> Params:
# def structured_drive_quantizer(rng: PRNGKey) -> aggregator.Aggregator:
#   def init():
#   def apply(
#       clients_params_and_weights: Iterable[Tuple[ClientId, Params, float]],
#       aggregator_state: CompressionState) -> Tuple[Params, CompressionState]:
#     def quantize_params_and_weight(client_id, params, weight):
# def terngrad_quantize(v: jnp.ndarray, rng: PRNGKey) -> jnp.ndarray:
# def terngrad_quantize_pytree(params: Params, rng: PRNGKey) -> Params:
# def terngrad_quantizer(rng: PRNGKey) -> aggregator.Aggregator:
#   def init():
#   def apply(
#       clients_params_and_weights: Iterable[Tuple[ClientId, Params, float]],
#       aggregator_state: CompressionState) -> Tuple[Params, CompressionState]:
#     def quantize_params_and_weight(client_params_and_weight, rng):


# Path: fedjax/aggregators/compression_test.py
from absl.testing import absltest
from fedjax.aggregators import compression
import jax
import jax.numpy as jnp
import numpy.testing as npt

    rng = jax.random.PRNGKey(42)
    compressed_sum = jnp.zeros_like(v)
    for _ in range(500):
      rng, use_rng = jax.random.split(rng)
      compressed_sum += compression.uniform_stochastic_quantize(v, 125, use_rng)
    npt.assert_array_almost_equal(compressed_sum / 500, v, decimal=2)

  def test_uniform_stochastic_quantize_unbiasedness_two_dim(self):
    v = jnp.array([[0., 1., 100.], [0.3, 2.3, 45.]])
    rng = jax.random.PRNGKey(42)
    compressed_sum = jnp.zeros_like(v)
    for _ in range(500):
      rng, use_rng = jax.random.split(rng)
      compressed_sum += compression.uniform_stochastic_quantize(v, 125, use_rng)
    npt.assert_array_almost_equal(compressed_sum / 500, v, decimal=1)

  def test_arithmetic_encoding_num_bits(self):
    v = jnp.array([1.0, 2.0, 3.0, 4.0, 5.0])
    npt.assert_array_almost_equal(
        compression.arithmetic_encoding_num_bits(v), [89.82311], decimal=3)

  def test_uniform_stochastic_quantizer(self):
    delta_params_and_weights = [('a', {
        'w': jnp.array([1., 2., 3.])
    }, 2.), ('b', {
        'w': jnp.array([2., 4., 6.])
    }, 4.), ('c', {
        'w': jnp.array([1., 3., 5.])
    }, 2.)]


    quantizer = compression.uniform_stochastic_quantizer(