# Path: t5/models/mesh_transformer.py
# DEPRECATED_GIN_REFERENCES = (
#     "configurable_vocabulary",
#     "get_sentencepiece_model_path",
#     "maybe_print_dataset",
#     "num_parallel_calls",
#     "SentencePieceVocabulary",
#     "t5.data.sentencepiece_vocabulary.SentencePieceVocabulary",
#     "t5.models.mesh_transformer.get_sentencepiece_model_path",
#     "train_model",
#     "vocabularies.Vocabulary",
#     "Vocabulary",
# )
# def mesh_train_dataset_fn(
#     mixture_or_task_name,
#     sequence_length,
#     vocabulary=None,
#     dataset_split=tfds.Split.TRAIN,
#     shuffle=True,
#     seed=None,
#     use_cached=False,
#     pack=True):
#   def _filter_features(ex):
# def mesh_inference_dataset_fn(
#     mixture_or_task_name,
#     sequence_length,
#     dataset_split,
#     shuffle=False,
#     seed=None,
#     vocabulary=None,
#     num_inference_examples=-1,
#     use_cached=False,
#     priming_sequence_length=None):
#   def _split_targets_for_primed_inference(ex):
#   def _prepare_for_unprimed_inference(ex):
#   def _get_dataset_for_single_task(task, sequence_length):
# def mesh_eval_dataset_fn(
#     mixture_or_task_name,
#     sequence_length,
#     dataset_split,
#     vocabulary=None,
#     num_eval_examples=-1,
#     use_cached=False,
#     pack=False,
#     shuffle_eval_examples=False,
#     seed=None):
#   def _get_dataset_for_single_task(task, sequence_length):
# def tsv_dataset_fn(
#     filename,
#     sequence_length,
#     dataset_split,
#     vocabulary,
#     shuffle_buffer_size=10000):
# def get_vocabulary(mixture_or_task_name=None):


from absl.testing import absltest
from t5.models import mesh_transformer
import seqio
import tensorflow.compat.v1 as tf
import tensorflow_datasets as tfds
# 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.models.mesh_transformer."""


tf.enable_eager_execution()


class MeshDatasetFnsTest(seqio.test_utils.FakeMixtureTest):

  def check_ds_shape(self, ds, sequence_length):
    for k, v in tf.data.get_output_shapes(ds).items():
      feat = k.split("_")[0]
      if len(v) == 0:  # pylint:disable=g-explicit-length-test
        expected_shape = []
      elif feat in sequence_length:
        expected_shape = [sequence_length[feat]]
      else:
        expected_shape = [None]
      self.assertEqual(expected_shape, v.as_list())

  def verify_mesh_dataset_fn(self, mixture_name, train, use_cached):
    if train:
      dataset_fn = mesh_transformer.mesh_train_dataset_fn
      split = tfds.Split.TRAIN
    else:

# Path: t5/models/mesh_transformer.py
# DEPRECATED_GIN_REFERENCES = (
#     "configurable_vocabulary",
#     "get_sentencepiece_model_path",
#     "maybe_print_dataset",
#     "num_parallel_calls",
#     "SentencePieceVocabulary",
#     "t5.data.sentencepiece_vocabulary.SentencePieceVocabulary",
#     "t5.models.mesh_transformer.get_sentencepiece_model_path",
#     "train_model",
#     "vocabularies.Vocabulary",
#     "Vocabulary",
# )
# def mesh_train_dataset_fn(
#     mixture_or_task_name,
#     sequence_length,
#     vocabulary=None,
#     dataset_split=tfds.Split.TRAIN,
#     shuffle=True,
#     seed=None,
#     use_cached=False,
#     pack=True):
#   def _filter_features(ex):
# def mesh_inference_dataset_fn(
#     mixture_or_task_name,
#     sequence_length,
#     dataset_split,
#     shuffle=False,
#     seed=None,
#     vocabulary=None,
#     num_inference_examples=-1,
#     use_cached=False,
#     priming_sequence_length=None):
#   def _split_targets_for_primed_inference(ex):
#   def _prepare_for_unprimed_inference(ex):
#   def _get_dataset_for_single_task(task, sequence_length):
# def mesh_eval_dataset_fn(
#     mixture_or_task_name,
#     sequence_length,
#     dataset_split,
#     vocabulary=None,
#     num_eval_examples=-1,
#     use_cached=False,
#     pack=False,
#     shuffle_eval_examples=False,
#     seed=None):
#   def _get_dataset_for_single_task(task, sequence_length):
# def tsv_dataset_fn(
#     filename,
#     sequence_length,
#     dataset_split,
#     vocabulary,
#     shuffle_buffer_size=10000):
# def get_vocabulary(mixture_or_task_name=None):


# Path: t5/models/mesh_transformer_test.py
from absl.testing import absltest
from t5.models import mesh_transformer
import seqio
import tensorflow.compat.v1 as tf
import tensorflow_datasets as tfds

# 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.models.mesh_transformer."""


tf.enable_eager_execution()


class MeshDatasetFnsTest(seqio.test_utils.FakeMixtureTest):

  def check_ds_shape(self, ds, sequence_length):
    for k, v in tf.data.get_output_shapes(ds).items():
      feat = k.split("_")[0]
      if len(v) == 0:  # pylint:disable=g-explicit-length-test
        expected_shape = []
      elif feat in sequence_length:
        expected_shape = [sequence_length[feat]]
      else:
        expected_shape = [None]
      self.assertEqual(expected_shape, v.as_list())

  def verify_mesh_dataset_fn(self, mixture_name, train, use_cached):
    if train:
      dataset_fn = mesh_transformer.mesh_train_dataset_fn
      split = tfds.Split.TRAIN
    else:

      dataset_fn = mesh_transformer.mesh_eval_dataset_fn