# 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
  vocab_size: int = 30522
  max_length: int = 512
  num_segments: int = 2

  bias_init: Initializer = nn.initializers.zeros
  kernel_init: Initializer = initializers.truncated_normal(stddev=0.02)

  layer_norm_epsilon: float = 1e-12
  dtype: jnp.dtype = jnp.float32
  # TODO: Set `hidden_activation` to jax.nn.gelu.

  dropout_rate: float = 0.1
  # TODO: Set `attention_probs_dropout_rate` to 0.1.


@dataclasses.dataclass
class BertLargeConfig(BertConfig):
  """BERT Large configuration."""

  hidden_size: int = 1024
  intermediate_dim: int = 4096

  num_hidden_layers: int = 24
  num_attention_heads: int = 16

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

  bias_init: Initializer = nn.initializers.zeros

# 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

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

  bias_init: Initializer = nn.initializers.zeros
  kernel_init: Initializer = initializers.truncated_normal(stddev=0.02)

  layer_norm_epsilon: float = 1e-12
  dtype: jnp.dtype = jnp.float32
  # TODO: Set `hidden_activation` to jax.nn.gelu.

  dropout_rate: float = 0.1
  # TODO: Set `attention_probs_dropout_rate` to 0.1.


@dataclasses.dataclass
class BertLargeConfig(BertConfig):
  """BERT Large configuration."""

  hidden_size: int = 1024
  intermediate_dim: int = 4096

  num_hidden_layers: int = 24
  num_attention_heads: int = 16

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

  bias_init: Initializer = nn.initializers.zeros

  kernel_init: Initializer = initializers.truncated_normal(stddev=0.02)