# 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):


from absl.testing import absltest
from jax import random
from flaxformer import testing_utils
from flaxformer.architectures.t5 import t5_architecture_test_utils as t5_test_utils
import numpy as np
# Copyright 2022 Google LLC.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# 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.

"""Tests for t5_architecture."""



expected_files = testing_utils.ExpectedJsonFiles(
    'flaxformer/architectures/t5/testdata')
check_params = expected_files.check_params_shapes_only


class ParallelFusedDecoderOnlyTest(absltest.TestCase):

  def test_decoder_shapes_fused_parallel(self):
    """Tests if the decoder parameter have the expected shapes."""

# 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/parallel_fused_decoder_test.py
from absl.testing import absltest
from jax import random
from flaxformer import testing_utils
from flaxformer.architectures.t5 import t5_architecture_test_utils as t5_test_utils
import numpy as np

# Copyright 2022 Google LLC.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# 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.

"""Tests for t5_architecture."""



expected_files = testing_utils.ExpectedJsonFiles(
    'flaxformer/architectures/t5/testdata')
check_params = expected_files.check_params_shapes_only


class ParallelFusedDecoderOnlyTest(absltest.TestCase):

  def test_decoder_shapes_fused_parallel(self):
    """Tests if the decoder parameter have the expected shapes."""

    decoder = t5_test_utils.make_parallel_fused_transformer_config()