class_docstr: Constrained MNIST dataset.
  
  Args:
      num_samples: number of examples per selected class/digit
      digits: list selected MNIST digits/classes
      kwargs: Same as MNIST
  
  Examples:
      >>> dataset = TrialMNIST(".", download=True)
      >>> len(dataset)
      300
      >>> sorted(set([d.item() for d in dataset.targets]))
      [0, 1, 2]
      >>> torch.bincount(dataset.targets)
      tensor([100, 100, 100])
class_docstr_tok: ["Constrained","MNIST","dataset",".","Args",":","num_samples",":","number","of","examples","per","selected","class\/digit","digits",":","list","selected","MNIST","digits\/classes","kwargs",":","Same","as","MNIST","Examples",":",">",">",">","dataset","=","TrialMNIST","(","``",".","``",",","download=True",")",">",">",">","len","(","dataset",")","300",">",">",">","sorted","(","set","(","[","d.item","(",")","for","d","in","dataset.targets","]",")",")","[","0",",","1",",","2","]",">",">",">","torch.bincount","(","dataset.targets",")","tensor","(","[","100",",","100",",","100","]",")"]
code: def _prepare_subset(
      full_data: Tensor,
      full_targets: Tensor,
      num_samples: int,
      digits: Sequence,
  ):
      classes = {d: 0 for d in digits}
      indexes = []
      for idx, target in enumerate(full_targets):
          label = target.item()
          if classes.get(label, float("inf")) >= num_samples:
              continue
          indexes.append(idx)
          classes[label] += 1
          if all(classes[k] >= num_samples for k in classes):
              break
      data = full_data[indexes]
      targets = full_targets[indexes]
      return data, targets
  
imports: import logging
  import os
  import random
  import time
  import urllib.request
  from typing import Optional, Sequence, Tuple
  import torch
  from torch import Tensor
  from torch.utils.data import Dataset
label_desc: Use image node_id 2 for calling the TrialMNIST obj's underlying member method code with example usage: obj._prepare_subset(full_data, full_targets, num_samples, digits) and returns: data,  targets