# Path: fedjax/aggregators/walsh_hadamard.py
# def walsh_hadamard_transform(
#     x: jnp.ndarray,
#     small_n: int = 2**7,
#     precision: Union[jax.lax.Precision, str] = 'highest') -> jnp.ndarray:
# def hadamard_matrix(n: int, dtype: jnp.dtype) -> jnp.ndarray:
# def structured_rotation(x: jnp.ndarray,
#                         rng: PRNGKey) -> Tuple[jnp.ndarray, jnp.ndarray]:
# def inverse_structured_rotation(x: jnp.ndarray, rng: PRNGKey,
#                                 original_shape: jnp.ndarray) -> jnp.ndarray:
# def structured_rotation_pytree(params: Params,
#                                rng: PRNGKey) -> Tuple[Params, Params]:
# def inverse_structured_rotation_pytree(params: Params, rng: PRNGKey,
#                                        shapes: Params) -> Params:


from absl.testing import absltest
from fedjax.aggregators import walsh_hadamard
import jax
import jax.numpy as jnp
import numpy.testing as npt
      for seed in range(5):
        n = 2**10
        x = jnp.array(jax.random.normal(jax.random.PRNGKey(seed), shape=[n]))
        expect = jnp.dot(
            walsh_hadamard.hadamard_matrix(n, x.dtype), x, precision=precision)
        for m in [2, 3, 4, 5, 6, 7, 8, 9, 10, 11]:
          small_n = 2**m
          with self.subTest(
              f'precision={precision}, seed={seed}, small_n={small_n}'):
            y = walsh_hadamard.walsh_hadamard_transform(x, small_n)
            self.assertEqual(x.shape, y.shape)
            self.assertEqual(x.dtype, y.dtype)
            npt.assert_allclose(y, expect, rtol=1e-4, atol=1e-4)

  def test_structured_rotation(self):
    n = 10
    m = 5
    x = jnp.array(jax.random.normal(jax.random.PRNGKey(100), shape=[m, n]))
    rng = jax.random.PRNGKey(10)
    y, x_shape = walsh_hadamard.structured_rotation(x, rng)
    self.assertEqual(x.shape, x_shape)
    z = walsh_hadamard.inverse_structured_rotation(y, rng, x_shape)
    npt.assert_allclose(z, x, rtol=1e-4, atol=1e-4)

  def test_structured_rotation_pytree(self):
    params = {
        'a': jnp.array([[1.0, 0.0, 0.0], [1.0, 2.0, 3.0]]),
        'b': jnp.array([[1.0, 0.0], [1.0, 2.0]])
    }
    rng = jax.random.PRNGKey(10)

# Path: fedjax/aggregators/walsh_hadamard.py
# def walsh_hadamard_transform(
#     x: jnp.ndarray,
#     small_n: int = 2**7,
#     precision: Union[jax.lax.Precision, str] = 'highest') -> jnp.ndarray:
# def hadamard_matrix(n: int, dtype: jnp.dtype) -> jnp.ndarray:
# def structured_rotation(x: jnp.ndarray,
#                         rng: PRNGKey) -> Tuple[jnp.ndarray, jnp.ndarray]:
# def inverse_structured_rotation(x: jnp.ndarray, rng: PRNGKey,
#                                 original_shape: jnp.ndarray) -> jnp.ndarray:
# def structured_rotation_pytree(params: Params,
#                                rng: PRNGKey) -> Tuple[Params, Params]:
# def inverse_structured_rotation_pytree(params: Params, rng: PRNGKey,
#                                        shapes: Params) -> Params:


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

      for seed in range(5):
        n = 2**10
        x = jnp.array(jax.random.normal(jax.random.PRNGKey(seed), shape=[n]))
        expect = jnp.dot(
            walsh_hadamard.hadamard_matrix(n, x.dtype), x, precision=precision)
        for m in [2, 3, 4, 5, 6, 7, 8, 9, 10, 11]:
          small_n = 2**m
          with self.subTest(
              f'precision={precision}, seed={seed}, small_n={small_n}'):
            y = walsh_hadamard.walsh_hadamard_transform(x, small_n)
            self.assertEqual(x.shape, y.shape)
            self.assertEqual(x.dtype, y.dtype)
            npt.assert_allclose(y, expect, rtol=1e-4, atol=1e-4)

  def test_structured_rotation(self):
    n = 10
    m = 5
    x = jnp.array(jax.random.normal(jax.random.PRNGKey(100), shape=[m, n]))
    rng = jax.random.PRNGKey(10)
    y, x_shape = walsh_hadamard.structured_rotation(x, rng)
    self.assertEqual(x.shape, x_shape)
    z = walsh_hadamard.inverse_structured_rotation(y, rng, x_shape)
    npt.assert_allclose(z, x, rtol=1e-4, atol=1e-4)

  def test_structured_rotation_pytree(self):
    params = {
        'a': jnp.array([[1.0, 0.0, 0.0], [1.0, 2.0, 3.0]]),
        'b': jnp.array([[1.0, 0.0], [1.0, 2.0]])
    }
    rng = jax.random.PRNGKey(10)

    y, shapes = walsh_hadamard.structured_rotation_pytree(params, rng)