# Path: meta_dataset/distribute_utils.py
# def _pad(tensor, total_size, position):
# def aggregate(tensor):


import collections
import tensorflow.compat.v1 as tf
from meta_dataset import distribute_utils
        'query_labels, support_class_ids, query_class_ids, way')):
  """Wraps an episode's data and facilitates creation of feed dict.

    This class provides the same functionality as an Episode, but it's intended
    for use in a distributed setting: it contains only a chunk of the support
    and query images in the episode (and their corresponding labels); basic
    accessor methods (support/query_images, support/query_labels,
    one_hot_support/query_labels, support/query_class_ids) will return only
    the local values.  The remaining properties, however, refer to the global
    episode stats (way, shots, unique_class_id's) and may be aggregated on
    demand.

    Args:
      support_images: A Tensor of images for fitting an episodic model.
      query_images: A Tensor of images for evaluating an episodic model.
      support_labels: A 1D Tensor, the matching support labels (numbers between
        0 and K-1, with K the number of classes involved in the episode).
      query_labels: A 1D Tensor, the matching query labels (numbers between 0
        and K-1, with K the number of classes involved in the episode).
      support_class_ids: A 1D Tensor, the matching support class ids (numbers
        between 0 and N-1, with N the number of classes in the full dataset).
      query_class_ids: A 1D Tensor, the matching query class ids (numbers
        between 0 and N-1, with N the number of classes in the full dataset).
      way: number of classes in the original episode.
  """

  @property
  def unique_class_ids(self):
    """Return global unique class id's for the episode."""
    return compute_unique_class_ids(

# Path: meta_dataset/distribute_utils.py
# def _pad(tensor, total_size, position):
# def aggregate(tensor):


# Path: meta_dataset/data/providers.py
import collections
import tensorflow.compat.v1 as tf
from meta_dataset import distribute_utils

        'query_labels, support_class_ids, query_class_ids, way')):
  """Wraps an episode's data and facilitates creation of feed dict.

    This class provides the same functionality as an Episode, but it's intended
    for use in a distributed setting: it contains only a chunk of the support
    and query images in the episode (and their corresponding labels); basic
    accessor methods (support/query_images, support/query_labels,
    one_hot_support/query_labels, support/query_class_ids) will return only
    the local values.  The remaining properties, however, refer to the global
    episode stats (way, shots, unique_class_id's) and may be aggregated on
    demand.

    Args:
      support_images: A Tensor of images for fitting an episodic model.
      query_images: A Tensor of images for evaluating an episodic model.
      support_labels: A 1D Tensor, the matching support labels (numbers between
        0 and K-1, with K the number of classes involved in the episode).
      query_labels: A 1D Tensor, the matching query labels (numbers between 0
        and K-1, with K the number of classes involved in the episode).
      support_class_ids: A 1D Tensor, the matching support class ids (numbers
        between 0 and N-1, with N the number of classes in the full dataset).
      query_class_ids: A 1D Tensor, the matching query class ids (numbers
        between 0 and N-1, with N the number of classes in the full dataset).
      way: number of classes in the original episode.
  """

  @property
  def unique_class_ids(self):
    """Return global unique class id's for the episode."""
    return compute_unique_class_ids(

        tf.concat((distribute_utils.aggregate(self.support_class_ids),