# Path: flaxformer/components/initializers.py
# def sinusoidal(min_scale: float = 1.0,
#                max_scale: float = 10000.0,
#                dtype: DType = jnp.float32) -> Initializer:
#   def init(key: PRNGKey, shape: Shape, dtype: DType = dtype) -> Array:
# def truncated_normal(mean: Union[float, Array] = 0.0,
#                      stddev: Union[float, Array] = 0.05,
#                      dtype: DType = jnp.float32) -> Initializer:
#   def init(key: PRNGKey, shape: Shape, dtype: DType = dtype) -> Array:


from absl.testing import absltest
from absl.testing import parameterized
from jax import numpy as jnp
from flaxformer.components import initializers
import jax
import numpy as np
            3.50468487e-01
        ],
        [
            -8.84301245e-01, -1.06949806e-01, 6.85548604e-01, 8.57080519e-01,
            -9.17811871e-01
        ],
        [
            -3.04965496e-01, -1.29926765e+00, -9.85570103e-02, -8.27740490e-01,
            -3.74757677e-01
        ],
    ]
    np.testing.assert_array_almost_equal(result, expected)

  def test_truncated_normal_returns_correct_mean(self):
    """Tests that truncated normal returns the requested mean."""
    mean = 2.0
    stddev = 1.0
    init_fn = initializers.truncated_normal(
        mean=mean, stddev=stddev, dtype=jnp.float32)
    rng = jax.random.PRNGKey(0)
    num_embeddings = 100
    features = 100
    shape = (num_embeddings, features)
    result = init_fn(rng, shape, jnp.float32)
    np.testing.assert_allclose(np.mean(result), mean, atol=0.01)

  def test_truncated_normal_returns_correct_stddev(self):
    """Tests that truncated normal returns the requested stddev."""
    mean = 0.0
    stddev = 0.05

# Path: flaxformer/components/initializers.py
# def sinusoidal(min_scale: float = 1.0,
#                max_scale: float = 10000.0,
#                dtype: DType = jnp.float32) -> Initializer:
#   def init(key: PRNGKey, shape: Shape, dtype: DType = dtype) -> Array:
# def truncated_normal(mean: Union[float, Array] = 0.0,
#                      stddev: Union[float, Array] = 0.05,
#                      dtype: DType = jnp.float32) -> Initializer:
#   def init(key: PRNGKey, shape: Shape, dtype: DType = dtype) -> Array:


# Path: flaxformer/components/initializers_test.py
from absl.testing import absltest
from absl.testing import parameterized
from jax import numpy as jnp
from flaxformer.components import initializers
import jax
import numpy as np

            3.50468487e-01
        ],
        [
            -8.84301245e-01, -1.06949806e-01, 6.85548604e-01, 8.57080519e-01,
            -9.17811871e-01
        ],
        [
            -3.04965496e-01, -1.29926765e+00, -9.85570103e-02, -8.27740490e-01,
            -3.74757677e-01
        ],
    ]
    np.testing.assert_array_almost_equal(result, expected)

  def test_truncated_normal_returns_correct_mean(self):
    """Tests that truncated normal returns the requested mean."""
    mean = 2.0
    stddev = 1.0
    init_fn = initializers.truncated_normal(
        mean=mean, stddev=stddev, dtype=jnp.float32)
    rng = jax.random.PRNGKey(0)
    num_embeddings = 100
    features = 100
    shape = (num_embeddings, features)
    result = init_fn(rng, shape, jnp.float32)
    np.testing.assert_allclose(np.mean(result), mean, atol=0.01)

  def test_truncated_normal_returns_correct_stddev(self):
    """Tests that truncated normal returns the requested stddev."""
    mean = 0.0
    stddev = 0.05

    init_fn = initializers.truncated_normal(