# 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
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""BERT Base configuration."""


@dataclasses.dataclass
class BertConfig(abc.ABC):
  """BERT configuration base dataclass."""
  # The size of embeddings/hidden layers, and the size of MLP intermediates.
  hidden_size: int
  intermediate_dim: int

  # The total number of layers and the number of attention heads in each layer.
  num_hidden_layers: int
  num_attention_heads: int

  # The size of the input/output vocabulary, the maximum supported length, and
  # the number of segments (type
  vocab_size: int
  max_length: int  # `max_position_embeddings` in legacy BertConfig.
  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).

# 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

#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""BERT Base configuration."""


@dataclasses.dataclass
class BertConfig(abc.ABC):
  """BERT configuration base dataclass."""
  # The size of embeddings/hidden layers, and the size of MLP intermediates.
  hidden_size: int
  intermediate_dim: int

  # The total number of layers and the number of attention heads in each layer.
  num_hidden_layers: int
  num_attention_heads: int

  # The size of the input/output vocabulary, the maximum supported length, and
  # the number of segments (type
  vocab_size: int
  max_length: int  # `max_position_embeddings` in legacy BertConfig.
  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