def test_get_dataset_enc_dec_sharded_and_packed(self):
    mixture_or_task_name = 'enc_dec_sharded_and_packed'
    x = [{'inputs': [7, 8], 'targets': [3, 9]}, {'inputs': [8, 4], 'targets': [4]}, {'inputs': [5, 6, 7], 'targets': [6]}]
    ds = create_default_dataset(x)
    dataset_fn = (lambda split, shuffle_files: ds)
    register_dummy_task(mixture_or_task_name, dataset_fn=dataset_fn)
    task_feature_lengths = {'inputs': 7, 'targets': 5}
    converter = feature_converters.EncDecFeatureConverter(pack=True)
    shard_info = dataset_providers.ShardInfo(index=0, num_shards=2)
    output_ds = dataset_providers.get_dataset(mixture_or_task_name=mixture_or_task_name, task_feature_lengths=task_feature_lengths, dataset_split='train', shuffle=False, feature_converter=converter, shard_info=shard_info)
    expected = {'encoder_input_tokens': [7, 8, 1, 5, 6, 7, 1], 'encoder_segment_ids': [1, 1, 1, 2, 2, 2, 2], 'encoder_positions': [0, 1, 2, 0, 1, 2, 3], 'decoder_target_tokens': [3, 9, 1, 6, 1], 'decoder_input_tokens': [0, 3, 9, 0, 6], 'decoder_loss_weights': [1, 1, 1, 1, 1], 'decoder_segment_ids': [1, 1, 1, 2, 2], 'decoder_positions': [0, 1, 2, 0, 1]}
    expected_dtypes = {feat: tf.int32 for feat in expected.keys()}
    assert_dataset(output_ds, expected, expected_dtypes=expected_dtypes)
def test_get_dataset_enc_dec_sharded_and_packed(self):
    mixture_or_task_name = 'enc_dec_sharded_and_packed'
    x = [{'inputs': [7, 8], 'targets': [3, 9]}, {'inputs': [8, 4], 'targets': [4]}, {'inputs': [5, 6, 7], 'targets': [6]}]
    ds = create_default_dataset(x)
    dataset_fn = (lambda split, shuffle_files: ds)
    register_dummy_task(mixture_or_task_name, dataset_fn=dataset_fn)
    task_feature_lengths = {'inputs': 7, 'targets': 5}
    converter = feature_converters.EncDecFeatureConverter(pack=True)
    shard_info = dataset_providers.ShardInfo(index=0, num_shards=2)
    output_ds = dataset_providers.get_dataset(mixture_or_task_name=mixture_or_task_name, task_feature_lengths=task_feature_lengths, dataset_split='train', shuffle=False, feature_converter=converter, shard_info=shard_info)
    expected = {'encoder_input_tokens': [7, 8, 1, 5, 6, 7, 1], 'encoder_segment_ids': [1, 1, 1, 2, 2, 2, 2], 'encoder_positions': [0, 1, 2, 0, 1, 2, 3], 'decoder_target_tokens': [3, 9, 1, 6, 1], 'decoder_input_tokens': [0, 3, 9, 0, 6], 'decoder_loss_weights': [1, 1, 1, 1, 1], 'decoder_segment_ids': [1, 1, 1, 2, 2], 'decoder_positions': [0, 1, 2, 0, 1]}
    expected_dtypes = {feat: tf.int32 for feat in expected.keys()}
    assert_dataset(output_ds, expected, expected_dtypes=expected_dtypes)