# 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/types.py



import abc
import dataclasses
from flax import linen as nn
from jax import numpy as jnp
from flaxformer.components import initializers
from flaxformer.types import Initializer  # pylint: disable=g-multiple-import
  num_segments: int  # `type_vocab_size` in legacy BertConfig.

  # Initializers, activations and dtypes for all the layers.
  # Legacy BertConfig has `initializer_range` which can be matched using
  # initializers.truncated_normal(stddev=initializer_range).
  bias_init: Initializer
  kernel_init: Initializer
  layer_norm_epsilon: float
  dtype: jnp.dtype
  # TODO: Support a `hidden_activation` config for the MLP.
  # `hidden_act` in legacy BertConfig.

  dropout_rate: float
  # TODO: Support a `attention_probs_dropout_rate` config.


@dataclasses.dataclass
class BertBaseConfig(BertConfig):
  """BERT Base configuration."""

  hidden_size: int = 768
  intermediate_dim: int = 3072

  num_hidden_layers: int = 12
  num_attention_heads: int = 12

  vocab_size: int = 30522
  max_length: int = 512
  num_segments: int = 2


# 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/types.py



# Path: flaxformer/architectures/bert/configs.py
import abc
import dataclasses
from flax import linen as nn
from jax import numpy as jnp
from flaxformer.components import initializers
from flaxformer.types import Initializer  # pylint: disable=g-multiple-import

  num_segments: int  # `type_vocab_size` in legacy BertConfig.

  # Initializers, activations and dtypes for all the layers.
  # Legacy BertConfig has `initializer_range` which can be matched using
  # initializers.truncated_normal(stddev=initializer_range).
  bias_init: Initializer
  kernel_init: Initializer
  layer_norm_epsilon: float
  dtype: jnp.dtype
  # TODO: Support a `hidden_activation` config for the MLP.
  # `hidden_act` in legacy BertConfig.

  dropout_rate: float
  # TODO: Support a `attention_probs_dropout_rate` config.


@dataclasses.dataclass
class BertBaseConfig(BertConfig):
  """BERT Base configuration."""

  hidden_size: int = 768
  intermediate_dim: int = 3072

  num_hidden_layers: int = 12
  num_attention_heads: int = 12

  vocab_size: int = 30522
  max_length: int = 512
  num_segments: int = 2


  bias_init: Initializer = nn.initializers.zeros