# Path: flaxformer/sharding.py
# class AxisNames(tuple):
# def axis_names(*names: str) -> partitioning.AxisMetadata:
# def reduce_fn(x, y):
# def _get_single_sowed_value(value) -> AxisNames:
# def get_axis_names(variables):
# def check_params_and_axis_names_match(variables):
#   def _flatten_with_joined_names(xs):
# 
# Path: flaxformer/testing_utils.py
# def param_shapes(params: Mapping[str, Any]) -> Dict[str, Any]:
# def param_dtypes_shapes(params: Mapping[str, Any]) -> Dict[str, Any]:
# def param_dtypes_shapes_axes(params: Mapping[str, Any],
#                              params_axes: Mapping[str, Any]) -> Dict[str, Any]:
#   def remove_axes_suffix(ks):
#   def _create_entry(param, param_axes):
# def format_params_shapes(params_shapes: Dict[str, Any]) -> str:
#   def re_compact_arrays(array_match) -> str:
#   def __init__(self, base_path: str):
#   def get_params(self, filename: str) -> Mapping[str, Any]:
#   def check_params(
#       self,
#       actual_params: Mapping[str, Any],
#       expected_filename: str,
#   ) -> None:
#   def check_params_and_axes(
#       self,
#       actual_params: Mapping[str, Any],
#       actual_params_axes: Mapping[str, Any],
#       expected_filename: str,
#   ) -> None:
#   def check_params_shapes_only(
#       self,
#       actual_params: Mapping[str, Any],
#       expected_filename: str,
#   ) -> None:
# class ExpectedJsonFiles:
# 
# Path: flaxformer/architectures/t5/t5_architecture_test_utils.py
# EMBEDDING_INIT = nn.initializers.normal(stddev=1.0)
# RELPOS_BIAS_INIT = nn.initializers.variance_scaling(1.0, 'fan_avg', 'uniform')
# ATTENTION_KERNEL_INIT = nn.initializers.variance_scaling(
#     1.0, 'fan_in', 'normal')
# MLP_KERNEL_INIT = nn.initializers.variance_scaling(1.0, 'fan_in',
#                                                    'truncated_normal')
# FINAL_KERNEL_INIT = nn.initializers.variance_scaling(1.0, 'fan_in',
#                                                      'truncated_normal')
# BIAS_INIT = nn.initializers.normal(stddev=1e-6)
# def make_token_emb1(vocab_size, dtype, features=13):
# def make_attention1(num_attn_heads, dtype):
# def make_mlp1(dtype):
# def _make_relative_position_bias(
#     num_attn_heads: int,
#     dtype: Any) -> relative_position_biases.RelativePositionBiases:
# def make_config1(scan_layers=False,
#                  layer_remat='legacy') -> t5_architecture.EncoderDecoder:
#   def _make_encoder_layer(shared_relative_position_bias):
#   def _make_decoder_layer(shared_relative_position_bias):
#   def _make_encoder(shared_token_embedder):
#   def _make_decoder(shared_token_embedder):
# def make_parallel_transformer_config() -> t5_architecture.EncoderDecoder:
#   def _make_encoder_layer(shared_relative_position_bias):
#   def _make_decoder_layer(shared_relative_position_bias):
#   def _make_encoder(shared_token_embedder):
#   def _make_decoder(shared_token_embedder):
# def make_parallel_fused_transformer_config() -> t5_architecture.DecoderOnly:
#   def _make_mq_attention(num_attn_heads, dtype):
#   def _make_fusion_mlp(dtype):
#   def _make_decoder_layer(shared_relative_position_bias):
#   def _make_output_logits():
#   def _embedder():
#   def _make_decoder(shared_token_embedder):
# def make_config2_shared_relative_position_bias(
# ) -> t5_architecture.EncoderDecoder:
#   def _make_encoder_layer(shared_relative_position_bias):
#   def _make_decoder_layer(shared_relative_position_bias):
#   def _make_encoder(*, shared_token_embedder=None):
#   def _make_decoder(*, shared_token_embedder=None):
# def make_config3_shared_token_embedder() -> t5_architecture.EncoderDecoder:
#   def _make_encoder_layer(shared_relative_position_bias):
#   def _make_decoder_layer(shared_relative_position_bias):
#   def _make_encoder(*, shared_token_embedder=None):
#   def _make_decoder(*, shared_token_embedder=None):
# def test_make_decoder_only1() -> t5_architecture.DecoderOnly:
#   def _make_decoder_layer(shared_relative_position_bias):
#   def make_output_logits():
#   def _make_decoder(*, shared_token_embedder=None):


import functools
import jax
import numpy as np
from absl.testing import absltest
from jax import random
from flaxformer import sharding
from flaxformer import testing_utils
from flaxformer.architectures.t5 import t5_architecture_test_utils as t5_test_utils
    np.testing.assert_allclose(output1, output3, rtol=1.5e-5)
    # Check scan_layers=True results
    np.testing.assert_allclose(output4, output5, rtol=1.5e-5)

  def test_scan_axis_annotations(self):
    encoder_input_tokens = np.array(
        [
            # Batch 1.
            [101, 183, 20, 75],
            # Batch 2.
            [101, 392, 19, 7],
        ],
        dtype=np.int32)
    transformer = t5_test_utils.make_config1(
        scan_layers=True, layer_remat='minimal')
    variables = transformer.init(
        random.PRNGKey(0),
        encoder_input_tokens,
        enable_dropout=False,
        method=transformer.encode,
    )

    # Check that the code can run when `params_axes` is not mutable too.
    transformer.apply(
        variables,
        encoder_input_tokens,
        enable_dropout=False,
        method=transformer.encode,
    )


# Path: flaxformer/sharding.py
# class AxisNames(tuple):
# def axis_names(*names: str) -> partitioning.AxisMetadata:
# def reduce_fn(x, y):
# def _get_single_sowed_value(value) -> AxisNames:
# def get_axis_names(variables):
# def check_params_and_axis_names_match(variables):
#   def _flatten_with_joined_names(xs):
# 
# Path: flaxformer/testing_utils.py
# def param_shapes(params: Mapping[str, Any]) -> Dict[str, Any]:
# def param_dtypes_shapes(params: Mapping[str, Any]) -> Dict[str, Any]:
# def param_dtypes_shapes_axes(params: Mapping[str, Any],
#                              params_axes: Mapping[str, Any]) -> Dict[str, Any]:
#   def remove_axes_suffix(ks):
#   def _create_entry(param, param_axes):
# def format_params_shapes(params_shapes: Dict[str, Any]) -> str:
#   def re_compact_arrays(array_match) -> str:
#   def __init__(self, base_path: str):
#   def get_params(self, filename: str) -> Mapping[str, Any]:
#   def check_params(
#       self,
#       actual_params: Mapping[str, Any],
#       expected_filename: str,
#   ) -> None:
#   def check_params_and_axes(
#       self,
#       actual_params: Mapping[str, Any],
#       actual_params_axes: Mapping[str, Any],
#       expected_filename: str,
#   ) -> None:
#   def check_params_shapes_only(
#       self,
#       actual_params: Mapping[str, Any],
#       expected_filename: str,
#   ) -> None:
# class ExpectedJsonFiles:
# 
# Path: flaxformer/architectures/t5/t5_architecture_test_utils.py
# EMBEDDING_INIT = nn.initializers.normal(stddev=1.0)
# RELPOS_BIAS_INIT = nn.initializers.variance_scaling(1.0, 'fan_avg', 'uniform')
# ATTENTION_KERNEL_INIT = nn.initializers.variance_scaling(
#     1.0, 'fan_in', 'normal')
# MLP_KERNEL_INIT = nn.initializers.variance_scaling(1.0, 'fan_in',
#                                                    'truncated_normal')
# FINAL_KERNEL_INIT = nn.initializers.variance_scaling(1.0, 'fan_in',
#                                                      'truncated_normal')
# BIAS_INIT = nn.initializers.normal(stddev=1e-6)
# def make_token_emb1(vocab_size, dtype, features=13):
# def make_attention1(num_attn_heads, dtype):
# def make_mlp1(dtype):
# def _make_relative_position_bias(
#     num_attn_heads: int,
#     dtype: Any) -> relative_position_biases.RelativePositionBiases:
# def make_config1(scan_layers=False,
#                  layer_remat='legacy') -> t5_architecture.EncoderDecoder:
#   def _make_encoder_layer(shared_relative_position_bias):
#   def _make_decoder_layer(shared_relative_position_bias):
#   def _make_encoder(shared_token_embedder):
#   def _make_decoder(shared_token_embedder):
# def make_parallel_transformer_config() -> t5_architecture.EncoderDecoder:
#   def _make_encoder_layer(shared_relative_position_bias):
#   def _make_decoder_layer(shared_relative_position_bias):
#   def _make_encoder(shared_token_embedder):
#   def _make_decoder(shared_token_embedder):
# def make_parallel_fused_transformer_config() -> t5_architecture.DecoderOnly:
#   def _make_mq_attention(num_attn_heads, dtype):
#   def _make_fusion_mlp(dtype):
#   def _make_decoder_layer(shared_relative_position_bias):
#   def _make_output_logits():
#   def _embedder():
#   def _make_decoder(shared_token_embedder):
# def make_config2_shared_relative_position_bias(
# ) -> t5_architecture.EncoderDecoder:
#   def _make_encoder_layer(shared_relative_position_bias):
#   def _make_decoder_layer(shared_relative_position_bias):
#   def _make_encoder(*, shared_token_embedder=None):
#   def _make_decoder(*, shared_token_embedder=None):
# def make_config3_shared_token_embedder() -> t5_architecture.EncoderDecoder:
#   def _make_encoder_layer(shared_relative_position_bias):
#   def _make_decoder_layer(shared_relative_position_bias):
#   def _make_encoder(*, shared_token_embedder=None):
#   def _make_decoder(*, shared_token_embedder=None):
# def test_make_decoder_only1() -> t5_architecture.DecoderOnly:
#   def _make_decoder_layer(shared_relative_position_bias):
#   def make_output_logits():
#   def _make_decoder(*, shared_token_embedder=None):


# Path: flaxformer/architectures/t5/t5_architecture_test.py
import functools
import jax
import numpy as np
from absl.testing import absltest
from jax import random
from flaxformer import sharding
from flaxformer import testing_utils
from flaxformer.architectures.t5 import t5_architecture_test_utils as t5_test_utils

    np.testing.assert_allclose(output1, output3, rtol=1.5e-5)
    # Check scan_layers=True results
    np.testing.assert_allclose(output4, output5, rtol=1.5e-5)

  def test_scan_axis_annotations(self):
    encoder_input_tokens = np.array(
        [
            # Batch 1.
            [101, 183, 20, 75],
            # Batch 2.
            [101, 392, 19, 7],
        ],
        dtype=np.int32)
    transformer = t5_test_utils.make_config1(
        scan_layers=True, layer_remat='minimal')
    variables = transformer.init(
        random.PRNGKey(0),
        encoder_input_tokens,
        enable_dropout=False,
        method=transformer.encode,
    )

    # Check that the code can run when `params_axes` is not mutable too.
    transformer.apply(
        variables,
        encoder_input_tokens,
        enable_dropout=False,
        method=transformer.encode,
    )


    sharding.check_params_and_axis_names_match(variables)