# 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
# 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 walsh_hadamard."""



class WalshHadamardTest(absltest.TestCase):

  def test_walsh_hadamard_transform(self):
    # TODO(wuke): Change bfloat16_3x to high once the PyPI JAX release catches
    # up.
    for precision in ['fastest', 'bfloat16_3x', 'highest']:
      for seed in range(5):
        n = 2**10
        x = jnp.array(jax.random.normal(jax.random.PRNGKey(seed), shape=[n]))
        expect = jnp.dot(

# 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

# 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 walsh_hadamard."""



class WalshHadamardTest(absltest.TestCase):

  def test_walsh_hadamard_transform(self):
    # TODO(wuke): Change bfloat16_3x to high once the PyPI JAX release catches
    # up.
    for precision in ['fastest', 'bfloat16_3x', 'highest']:
      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)