# Path: flaxformer/components/dense.py
# def dense(features, name, inputs, kernel_axis_names):
#   if self.weight_params is not None or self.act_params is not None:
#     # TODO: Push the "quantized vs not" decision down into the
#     # AQT library. Currently we make that decision here, because the AQT
#     # library doesn't support DenseGeneral, so there's extra reshapes here
#     # whose performance impact I don't know.
#     aqt_context = aqt_config.QuantContext(
#         update_bounds=False, collect_acts_stats=False)
#     aqt_hparams = aqt_flax_layers.DenseAqt.HParams(
#         weight_prec=self.weight_params.prec,
#         weight_half_shift=self.weight_params.half_shift,
#         quant_act=self.act_params,  # currently supports fixed bounds only.
#         quant_type=aqt.QuantType.aqt,
#         weight_quant_granularity=aqt_config.QuantGranularity.per_channel,
#     )
#     batch, seq_len, channels = inputs.shape
#     inputs = inputs.reshape((batch * seq_len, channels))
#     result = aqt_flax_layers.DenseAqt(
#         features=features,
#         hparams=aqt_hparams,
#         train=enable_dropout,
#         quant_context=aqt_context,
#         paxis_name=None,
#         # No "cross-replica" reduction expressed in the XLA graph at this
#         # stage. Will be imposed later, automatically, by XLA SPMD.
#         use_bias=self.use_bias,
#         kernel_init=self.kernel_init,
#         bias_init=self.bias_init,
#         dtype=self.dtype,
#         kernel_axis_names=kernel_axis_names,
#         name=name,
#     )(inputs, padding_mask=None)
#     return result.reshape((batch, seq_len, features))
#   else:
#     return DenseGeneral(
#         features,
#         use_bias=self.use_bias,
#         dtype=self.dtype,
#         kernel_init=self.kernel_init,
#         bias_init=self.bias_init,
#         kernel_axis_names=kernel_axis_names,
#         name=name)(
#             inputs)
# 
# 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/types.py

# 
# Path: flaxformer/types.py

# 
# Path: flaxformer/types.py



from typing import Optional
from flax import linen as nn
from flax.linen.initializers import zeros
from flaxformer.components import dense
from flaxformer.components import initializers
from flaxformer.types import Activation
from flaxformer.types import Array
from flaxformer.types import DType
from flaxformer.types import Initializer
import jax.numpy as jnp
    self.dense = dense.DenseGeneral(
        features=self.num_classes,
        use_bias=self.use_bias,
        dtype=self.dtype,
        kernel_init=self.kernel_init,
        name='dense')

  def __call__(self,
               encoded_inputs: Array,
               *,
               enable_dropout: Optional[bool] = None,
               **unused_kwargs) -> Array:
    """Pools the CLS emebdding and projects into the logits.

    Args:
      encoded_inputs: The inputs (e.g., token's embeddings) that come from the
        final layer of the BERT encoder. <float32>[batch_size, seq_length,
        hidden_size].
      enable_dropout: Enables dropout when set to True.
      **unused_kwargs: unused.

    Returns:
      An array of logits <float32>[batch_size, num_classes].
    """
    if enable_dropout is not None:
      deterministic = not enable_dropout
    else:
      deterministic = None
    cls_embedding = self.pooler(encoded_inputs)
    cls_embedding = self.dropout_layer(cls_embedding, deterministic)

# Path: flaxformer/components/dense.py
# def dense(features, name, inputs, kernel_axis_names):
#   if self.weight_params is not None or self.act_params is not None:
#     # TODO: Push the "quantized vs not" decision down into the
#     # AQT library. Currently we make that decision here, because the AQT
#     # library doesn't support DenseGeneral, so there's extra reshapes here
#     # whose performance impact I don't know.
#     aqt_context = aqt_config.QuantContext(
#         update_bounds=False, collect_acts_stats=False)
#     aqt_hparams = aqt_flax_layers.DenseAqt.HParams(
#         weight_prec=self.weight_params.prec,
#         weight_half_shift=self.weight_params.half_shift,
#         quant_act=self.act_params,  # currently supports fixed bounds only.
#         quant_type=aqt.QuantType.aqt,
#         weight_quant_granularity=aqt_config.QuantGranularity.per_channel,
#     )
#     batch, seq_len, channels = inputs.shape
#     inputs = inputs.reshape((batch * seq_len, channels))
#     result = aqt_flax_layers.DenseAqt(
#         features=features,
#         hparams=aqt_hparams,
#         train=enable_dropout,
#         quant_context=aqt_context,
#         paxis_name=None,
#         # No "cross-replica" reduction expressed in the XLA graph at this
#         # stage. Will be imposed later, automatically, by XLA SPMD.
#         use_bias=self.use_bias,
#         kernel_init=self.kernel_init,
#         bias_init=self.bias_init,
#         dtype=self.dtype,
#         kernel_axis_names=kernel_axis_names,
#         name=name,
#     )(inputs, padding_mask=None)
#     return result.reshape((batch, seq_len, features))
#   else:
#     return DenseGeneral(
#         features,
#         use_bias=self.use_bias,
#         dtype=self.dtype,
#         kernel_init=self.kernel_init,
#         bias_init=self.bias_init,
#         kernel_axis_names=kernel_axis_names,
#         name=name)(
#             inputs)
# 
# 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/types.py

# 
# Path: flaxformer/types.py

# 
# Path: flaxformer/types.py



# Path: flaxformer/architectures/bert/heads.py
from typing import Optional
from flax import linen as nn
from flax.linen.initializers import zeros
from flaxformer.components import dense
from flaxformer.components import initializers
from flaxformer.types import Activation
from flaxformer.types import Array
from flaxformer.types import DType
from flaxformer.types import Initializer
import jax.numpy as jnp

    self.dense = dense.DenseGeneral(
        features=self.num_classes,
        use_bias=self.use_bias,
        dtype=self.dtype,
        kernel_init=self.kernel_init,
        name='dense')

  def __call__(self,
               encoded_inputs: Array,
               *,
               enable_dropout: Optional[bool] = None,
               **unused_kwargs) -> Array:
    """Pools the CLS emebdding and projects into the logits.

    Args:
      encoded_inputs: The inputs (e.g., token's embeddings) that come from the
        final layer of the BERT encoder. <float32>[batch_size, seq_length,
        hidden_size].
      enable_dropout: Enables dropout when set to True.
      **unused_kwargs: unused.

    Returns:
      An array of logits <float32>[batch_size, num_classes].
    """
    if enable_dropout is not None:
      deterministic = not enable_dropout
    else:
      deterministic = None
    cls_embedding = self.pooler(encoded_inputs)
    cls_embedding = self.dropout_layer(cls_embedding, deterministic)

    logits = self.dense(cls_embedding)