# Path: fedjax/core/typing.py



import math
import jax
import jax.numpy as jnp
import scipy
from typing import Tuple, Union
from fedjax.core.typing import PRNGKey, Params
  d = 2**math.ceil(math.log2(x_flat.size))
  w = jnp.pad(x_flat, (0, d - x.size))
  rademacher = jax.random.rademacher(rng, w.shape)
  return walsh_hadamard_transform(w * rademacher) / jnp.sqrt(d), jnp.array(
      x.shape)


def inverse_structured_rotation(x: jnp.ndarray, rng: PRNGKey,
                                original_shape: jnp.ndarray) -> jnp.ndarray:
  """Computes (HD)^(-1)(x)/sqrt(d).

  Here where H is the walsh Hadamard matrix, d is the dimensionlity of x, and D
  is a random Rademacher matrix.

  Args:
    x: rotated array, which needs to be unrotated.
    rng: PRNGKey used for rotation.
    original_shape: desired shape of the output.

  Returns:
    Output of (HD)^(-1)(x)/sqrt(d) with appropriate shape.
  """
  rademacher = jax.random.rademacher(rng, x.shape)
  w = walsh_hadamard_transform(x) * rademacher / jnp.sqrt(x.size)
  original_size = jnp.prod(original_shape)
  y_flat = w.take(jnp.arange(original_size))
  return jnp.reshape(y_flat, original_shape)


def structured_rotation_pytree(params: Params,

# Path: fedjax/core/typing.py



# Path: fedjax/aggregators/walsh_hadamard.py
import math
import jax
import jax.numpy as jnp
import scipy
from typing import Tuple, Union
from fedjax.core.typing import PRNGKey, Params

  d = 2**math.ceil(math.log2(x_flat.size))
  w = jnp.pad(x_flat, (0, d - x.size))
  rademacher = jax.random.rademacher(rng, w.shape)
  return walsh_hadamard_transform(w * rademacher) / jnp.sqrt(d), jnp.array(
      x.shape)


def inverse_structured_rotation(x: jnp.ndarray, rng: PRNGKey,
                                original_shape: jnp.ndarray) -> jnp.ndarray:
  """Computes (HD)^(-1)(x)/sqrt(d).

  Here where H is the walsh Hadamard matrix, d is the dimensionlity of x, and D
  is a random Rademacher matrix.

  Args:
    x: rotated array, which needs to be unrotated.
    rng: PRNGKey used for rotation.
    original_shape: desired shape of the output.

  Returns:
    Output of (HD)^(-1)(x)/sqrt(d) with appropriate shape.
  """
  rademacher = jax.random.rademacher(rng, x.shape)
  w = walsh_hadamard_transform(x) * rademacher / jnp.sqrt(x.size)
  original_size = jnp.prod(original_shape)
  y_flat = w.take(jnp.arange(original_size))
  return jnp.reshape(y_flat, original_shape)


def structured_rotation_pytree(params: Params,

                               rng: PRNGKey) -> Tuple[Params, Params]: