# Path: flaxformer/types.py

# 
# Path: flaxformer/types.py

# 
# Path: flaxformer/types.py

# 
# Path: flaxformer/types.py

# 
# Path: flaxformer/types.py



from typing import Union
from jax import numpy as jnp
from flaxformer.types import Array
from flaxformer.types import DType
from flaxformer.types import Initializer
from flaxformer.types import PRNGKey
from flaxformer.types import Shape
import jax
import numpy as np
    min_scale: Minimum frequency-scale in sine grating.
    max_scale: Maximum frequency-scale in sine grating.
    dtype: The DType of the returned values.

  Returns:
    The sinusoidal initialization function.
  """

  def init(key: PRNGKey, shape: Shape, dtype: DType = dtype) -> Array:
    """Sinusoidal init."""
    del key
    if dtype != np.float32:
      raise ValueError('The sinusoidal initializer only supports float32.')
    if len(list(shape)) != 2:
      raise ValueError(
          f'Expected a 2D shape (max_len, features), but got {shape}.')
    max_len, features = shape
    pe = np.zeros((max_len, features), dtype=dtype)
    position = np.arange(0, max_len)[:, np.newaxis]
    scale_factor = -np.log(max_scale / min_scale) / (features // 2 - 1)
    div_term = min_scale * np.exp(np.arange(0, features // 2) * scale_factor)
    pe[:, :features // 2] = np.sin(position * div_term)
    pe[:, features // 2:2 * (features // 2)] = np.cos(position * div_term)
    return jnp.array(pe)

  return init


def truncated_normal(mean: Union[float, Array] = 0.0,
                     stddev: Union[float, Array] = 0.05,

# Path: flaxformer/types.py

# 
# Path: flaxformer/types.py

# 
# Path: flaxformer/types.py

# 
# Path: flaxformer/types.py

# 
# Path: flaxformer/types.py



# Path: flaxformer/components/initializers.py
from typing import Union
from jax import numpy as jnp
from flaxformer.types import Array
from flaxformer.types import DType
from flaxformer.types import Initializer
from flaxformer.types import PRNGKey
from flaxformer.types import Shape
import jax
import numpy as np

    min_scale: Minimum frequency-scale in sine grating.
    max_scale: Maximum frequency-scale in sine grating.
    dtype: The DType of the returned values.

  Returns:
    The sinusoidal initialization function.
  """

  def init(key: PRNGKey, shape: Shape, dtype: DType = dtype) -> Array:
    """Sinusoidal init."""
    del key
    if dtype != np.float32:
      raise ValueError('The sinusoidal initializer only supports float32.')
    if len(list(shape)) != 2:
      raise ValueError(
          f'Expected a 2D shape (max_len, features), but got {shape}.')
    max_len, features = shape
    pe = np.zeros((max_len, features), dtype=dtype)
    position = np.arange(0, max_len)[:, np.newaxis]
    scale_factor = -np.log(max_scale / min_scale) / (features // 2 - 1)
    div_term = min_scale * np.exp(np.arange(0, features // 2) * scale_factor)
    pe[:, :features // 2] = np.sin(position * div_term)
    pe[:, features // 2:2 * (features // 2)] = np.cos(position * div_term)
    return jnp.array(pe)

  return init


def truncated_normal(mean: Union[float, Array] = 0.0,
                     stddev: Union[float, Array] = 0.05,

                     dtype: DType = jnp.float32) -> Initializer: