# Path: mesh_tensorflow/transformer/dataset.py
# def pack_or_pad(
#     dataset, length, pack=True, feature_keys=None, ensure_eos=False):
# def ensure_dataset_eos(dataset, feature_keys=None):
#   def _ensure_eos(k, v):
# def encode_dataset(dataset, vocabulary):
#   def encode(features):
# def pretokenized_tfds_dataset(dataset_name=gin.REQUIRED,
#                               text2self=gin.REQUIRED,
#                               tfds_data_dir=gin.REQUIRED,
#                               dataset_split=gin.REQUIRED,
#                               batch_size=None,
#                               sequence_length=gin.REQUIRED,
#                               vocabulary=None):
#   def shift_and_append_eos(t):
#   def feature_map(inputs, targets):
# def sample_from_text_line_datasets(glob_weight_list,
#                                    shuffle_buffer_size=100000,
#                                    prefetch=1000):  # pylint: disable=missing-docstring
# def make_text_line_dataset(glob=gin.REQUIRED):
# def simple_text_line_dataset(glob=gin.REQUIRED, shuffle_buffer_size=100000):
# def packed_parallel_tsv_dataset(dataset=gin.REQUIRED,
#                                 dataset_split=gin.REQUIRED,
#                                 batch_size=None,
#                                 sequence_length=gin.REQUIRED,
#                                 vocabulary=gin.REQUIRED,
#                                 append_eos=True,
#                                 eos_id=1,
#                                 max_encoded_len=0):
#   def _parse_fn(record):  # pylint: disable=missing-docstring
#   def _encode_fn(features):  # pylint: disable=missing-docstring
#   def _filter_fn(features):  # pylint: disable=missing-docstring
# def untokenized_tfds_dataset(dataset_name=gin.REQUIRED,
#                              text2self=gin.REQUIRED,
#                              tfds_data_dir=gin.REQUIRED,
#                              dataset_split=gin.REQUIRED,
#                              batch_size=None,
#                              sequence_length=gin.REQUIRED,
#                              vocabulary=gin.REQUIRED,
#                              pack=gin.REQUIRED):
# def supervised_to_dict(dataset, text2self):
#   def my_fn(inputs, targets):
# def encode_all_features(dataset, vocabulary):
#   def my_fn(features):
# def pretokenized_tfrecord_dataset(filenames,
#                                   text2self,
#                                   eos_included,
#                                   repeat,
#                                   batch_size,
#                                   sequence_length,
#                                   vocab_shift=0):
#   def decode_example(serialized_example):
# def pretokenized_t2t_dataset(dataset_name=gin.REQUIRED,
#                              text2self=False,
#                              data_dir=gin.REQUIRED,
#                              dataset_split="train",
#                              batch_size=None,
#                              sequence_length=gin.REQUIRED,
#                              vocabulary=None,
#                              eos_included=True,
#                              vocab_shift=0):
# def pack_dataset(dataset, length, keys=None, use_custom_ops=False):
#   def my_fn(x):
# def _pack_with_tf_ops(dataset, keys, length):
#   def write_packed_example(partial, outputs):
#   def map_fn(x):
#     def cond_fn(i, partial, outputs):
#     def body_fn(i, partial, outputs):
#       def false_fn():
#       def true_fn():
# def _pack_with_custom_ops(dataset, keys, length):
#   def custom_pack_batch(x):
# def trim_and_pad_dataset(dataset, length, feature_keys=None):
#   def _trim_and_pad(k, t):
# def pad_dataset_with_zeroed_out_examples(ds):
#   def _zero_out(x):


from absl.testing import absltest
from absl.testing import parameterized
from mesh_tensorflow.transformer import dataset
import numpy as np
import tensorflow.compat.v1 as tf
import tensorflow_datasets as tfds
        if k in expected_dtypes:
          self.assertEqual(v.dtype.type, expected_dtypes[k])

  @parameterized.parameters(*_PACK_PARAMETERS)
  def test_pack_dataset(self, use_custom_ops):
    x = [{"inputs": [7, 8, 5, 1], "targets": [3, 9, 1], "idx": [0]},
         {"inputs": [8, 4, 9, 3, 1], "targets": [4, 1], "idx": [1]}]
    ds = create_default_dataset(x, feature_names=("inputs", "targets", "idx"))
    packed_ds = dataset.pack_dataset(
        ds,
        length={"inputs": 10, "targets": 7},
        keys=("inputs", "targets"),
        use_custom_ops=use_custom_ops)

    expected = [{
        "inputs": [7, 8, 5, 1, 8, 4, 9, 3, 1, 0],
        "inputs_segmentation": [1, 1, 1, 1, 2, 2, 2, 2, 2, 0],
        "inputs_position": [0, 1, 2, 3, 0, 1, 2, 3, 4, 0],
        "targets": [3, 9, 1, 4, 1, 0, 0],
        "targets_position": [0, 1, 2, 0, 1, 0, 0],
        "targets_segmentation": [1, 1, 1, 2, 2, 0, 0],
    }]
    self.assert_dataset(
        packed_ds, expected, {"inputs": tf.int32, "targets": tf.int32})

  @parameterized.parameters(*_PACK_PARAMETERS)
  def test_pack_dataset_no_eos(self, use_custom_ops):
    x = [{"inputs": [7, 8, 5], "targets": [3, 9]},
         {"inputs": [8, 4, 9, 3], "targets": [4]}]
    ds = create_default_dataset(x)

# Path: mesh_tensorflow/transformer/dataset.py
# def pack_or_pad(
#     dataset, length, pack=True, feature_keys=None, ensure_eos=False):
# def ensure_dataset_eos(dataset, feature_keys=None):
#   def _ensure_eos(k, v):
# def encode_dataset(dataset, vocabulary):
#   def encode(features):
# def pretokenized_tfds_dataset(dataset_name=gin.REQUIRED,
#                               text2self=gin.REQUIRED,
#                               tfds_data_dir=gin.REQUIRED,
#                               dataset_split=gin.REQUIRED,
#                               batch_size=None,
#                               sequence_length=gin.REQUIRED,
#                               vocabulary=None):
#   def shift_and_append_eos(t):
#   def feature_map(inputs, targets):
# def sample_from_text_line_datasets(glob_weight_list,
#                                    shuffle_buffer_size=100000,
#                                    prefetch=1000):  # pylint: disable=missing-docstring
# def make_text_line_dataset(glob=gin.REQUIRED):
# def simple_text_line_dataset(glob=gin.REQUIRED, shuffle_buffer_size=100000):
# def packed_parallel_tsv_dataset(dataset=gin.REQUIRED,
#                                 dataset_split=gin.REQUIRED,
#                                 batch_size=None,
#                                 sequence_length=gin.REQUIRED,
#                                 vocabulary=gin.REQUIRED,
#                                 append_eos=True,
#                                 eos_id=1,
#                                 max_encoded_len=0):
#   def _parse_fn(record):  # pylint: disable=missing-docstring
#   def _encode_fn(features):  # pylint: disable=missing-docstring
#   def _filter_fn(features):  # pylint: disable=missing-docstring
# def untokenized_tfds_dataset(dataset_name=gin.REQUIRED,
#                              text2self=gin.REQUIRED,
#                              tfds_data_dir=gin.REQUIRED,
#                              dataset_split=gin.REQUIRED,
#                              batch_size=None,
#                              sequence_length=gin.REQUIRED,
#                              vocabulary=gin.REQUIRED,
#                              pack=gin.REQUIRED):
# def supervised_to_dict(dataset, text2self):
#   def my_fn(inputs, targets):
# def encode_all_features(dataset, vocabulary):
#   def my_fn(features):
# def pretokenized_tfrecord_dataset(filenames,
#                                   text2self,
#                                   eos_included,
#                                   repeat,
#                                   batch_size,
#                                   sequence_length,
#                                   vocab_shift=0):
#   def decode_example(serialized_example):
# def pretokenized_t2t_dataset(dataset_name=gin.REQUIRED,
#                              text2self=False,
#                              data_dir=gin.REQUIRED,
#                              dataset_split="train",
#                              batch_size=None,
#                              sequence_length=gin.REQUIRED,
#                              vocabulary=None,
#                              eos_included=True,
#                              vocab_shift=0):
# def pack_dataset(dataset, length, keys=None, use_custom_ops=False):
#   def my_fn(x):
# def _pack_with_tf_ops(dataset, keys, length):
#   def write_packed_example(partial, outputs):
#   def map_fn(x):
#     def cond_fn(i, partial, outputs):
#     def body_fn(i, partial, outputs):
#       def false_fn():
#       def true_fn():
# def _pack_with_custom_ops(dataset, keys, length):
#   def custom_pack_batch(x):
# def trim_and_pad_dataset(dataset, length, feature_keys=None):
#   def _trim_and_pad(k, t):
# def pad_dataset_with_zeroed_out_examples(ds):
#   def _zero_out(x):


# Path: mesh_tensorflow/transformer/dataset_test.py
from absl.testing import absltest
from absl.testing import parameterized
from mesh_tensorflow.transformer import dataset
import numpy as np
import tensorflow.compat.v1 as tf
import tensorflow_datasets as tfds

        if k in expected_dtypes:
          self.assertEqual(v.dtype.type, expected_dtypes[k])

  @parameterized.parameters(*_PACK_PARAMETERS)
  def test_pack_dataset(self, use_custom_ops):
    x = [{"inputs": [7, 8, 5, 1], "targets": [3, 9, 1], "idx": [0]},
         {"inputs": [8, 4, 9, 3, 1], "targets": [4, 1], "idx": [1]}]
    ds = create_default_dataset(x, feature_names=("inputs", "targets", "idx"))
    packed_ds = dataset.pack_dataset(
        ds,
        length={"inputs": 10, "targets": 7},
        keys=("inputs", "targets"),
        use_custom_ops=use_custom_ops)

    expected = [{
        "inputs": [7, 8, 5, 1, 8, 4, 9, 3, 1, 0],
        "inputs_segmentation": [1, 1, 1, 1, 2, 2, 2, 2, 2, 0],
        "inputs_position": [0, 1, 2, 3, 0, 1, 2, 3, 4, 0],
        "targets": [3, 9, 1, 4, 1, 0, 0],
        "targets_position": [0, 1, 2, 0, 1, 0, 0],
        "targets_segmentation": [1, 1, 1, 2, 2, 0, 0],
    }]
    self.assert_dataset(
        packed_ds, expected, {"inputs": tf.int32, "targets": tf.int32})

  @parameterized.parameters(*_PACK_PARAMETERS)
  def test_pack_dataset_no_eos(self, use_custom_ops):
    x = [{"inputs": [7, 8, 5], "targets": [3, 9]},
         {"inputs": [8, 4, 9, 3], "targets": [4]}]
    ds = create_default_dataset(x)

    packed_ds = dataset.pack_dataset(