text: import torch
  import torchvision
  
  from torch.utils.data import Dataset, Subset
  from torchvision.transforms import ToTensor
  
  from src.data.cmnist_dist import make_joint_distribution, CMNIST_NAME, CMNIST_VERSIONS
  
  CMNIST_SIZE = 28 ** 2 * 2
  CMNIST_SHAPE = [2, 28, 28]
  CMNIST_N_CLASSES = 2
  CMNIST_N_ENVS = 2
  
  MNIST_TRAIN = 'train'
  MNIST_VALID = 'valid'
  MNIST_TEST = 'test'
  MNIST_TRAIN_VALID = 'train+valid'
  MNIST_TRAIN_SPLITS = [MNIST_TRAIN, MNIST_VALID, MNIST_TRAIN_VALID]
  MNIST_SPLITS = MNIST_TRAIN_SPLITS + [MNIST_TEST]
  MNIST_TRAIN_EXAMPLES = 50000
  
  
  # Wrapper for the torchvision MNIST dataset with validation split
  class MNIST(Dataset):
      def __init__(self, root, split, keep_in_memory=False, device='cpu', **params):
          super(MNIST, self).__init__()
  
          dataset = torchvision.datasets.MNIST(root=root,
                                               train=split in MNIST_TRAIN_SPLITS,
                                               transform=ToTensor(), **params)
  
          if split == MNIST_TRAIN:
              dataset = Subset(dataset, range(MNIST_TRAIN_EXAMPLES))
          elif split == MNIST_VALID:
              dataset = Subset(dataset, range(MNIST_TRAIN_EXAMPLES, len(dataset)))
          elif not (split == MNIST_TEST) and not (split == MNIST_TRAIN_VALID):
              raise Exception('The possible splits are %s' % ', '.join(MNIST_SPLITS))
          self.dataset = dataset
          self.keep_in_memory = keep_in_memory
          if keep_in_memory:
              cache = {
                  'x': torch.FloatTensor(len(dataset), 1, 28, 28).to(device),
                  'y': torch.LongTensor(len(dataset)).to(device)
              }
          else:
              cache = None
  
          self.cache = cache
          self.device = device
          self.stored_ids = []
  
      def __getitem__(self, index):
          if not self.keep_in_memory or not (index in self.stored_ids):
              x, y = self.dataset[index]
              if self.keep_in_memory:
                  self.cache['x'][index] = x.to(self.device)
                  self.cache['y'][index] = y
                  self.stored_ids.append(index)
          else:
              x = self.cache['x'][index]
              y = self.cache['y'][index]
          return {'x': x, 'y': torch.LongTensor([y])}
  
      def __len__(self):
          return len(self.dataset)
  
  
  # Implementation of the CMNIST, d-CMNIST and y-CMNIST datasets for pytorch
  class CMNIST(MNIST):
      def __init__(self, root, version=CMNIST_NAME, sample_once=False, t=1, **params):
          super(CMNIST, self).__init__(root=root, **params)
  
          assert version in CMNIST_VERSIONS
          assert t in [0, 1]
  
          self.dist = make_joint_distribution(version).condition_on('t', t)
          self.sample_once = sample_once
          self.sampled_data = {}
  
      def __getitem__(self, index):
          if index in self.sampled_data:
              data = self.sampled_data[index]
          else:
              data = super(CMNIST, self).__getitem__(index)
              x = data['x']
              d = data['y']
  
              # sample from p(e,y,c|d) to determine color, label and environment
              sample = self.dist.condition_on('d', d).sample()
  
              # Concatenate an empty channel (red)
              x = torch.cat([x, x * 0], 0)
  
              # If the color is 1, make the empty channel the first (green)
              if sample['c'] == 1:
                  x = torch.roll(x, 1, 0)
  
              data = {'x': x, 'y': sample['y'], 'e': sample['e']}
          if self.sample_once:
              self.sampled_data[index] = data
          return data
  