code: def _split_tcga(tcga_metadataset, counts):
      all_allowed_samples = tcga_metadataset.task_ids
  
      # We first uniquely assing every sample to a task
      sample_to_task_assignment = _assign_samples(tcga_metadataset)
  
      keys = [i for i in all_allowed_samples.keys()]
      difference = set(sample_to_task_assignment.keys()).difference(
          set(keys)
      )
  
      unassigned_samples = OrderedSet()
      for key in difference:
          unassigned_samples.update(sample_to_task_assignment[key])
  
      # Second we split the metadataset
      # with a torch-based random sample
      permutation = torch.randperm(len(keys)).numpy()
  
      metadatasets = []
      start = 0
      end = 0
      for count in counts:
          end += count
          current_keys = [
              keys[index] for index in permutation[start:end]
          ]
          metadatasets.append(
              {
                  key: sample_to_task_assignment[key]
                  for key in current_keys
              }
          )
          start = end
  
      expanded_metadatasets = [None] * len(metadatasets)
      order = np.argsort(
          [len(metadataset) for metadataset in metadatasets]
      )
  
      # Finally we expand the tasks by reusing samples wherever possible in the sets
      blacklist = OrderedSet()
      for i in order:
          additional_samples = unassigned_samples.difference(blacklist)
          (
              expanded_metadataset,
              used_additional_samples,
          ) = _expand_sample_usage(
              metadatasets[i], all_allowed_samples, additional_samples
          )
          expanded_metadatasets[i] = expanded_metadataset
          blacklist.update(used_additional_samples)
  
      tcga_metadatasets = []
      tcga_metadataset.close()
      preloaded = tcga_metadataset.preloaded
      for metadataset in expanded_metadatasets:
          current_tcga_metadataset = copy.deepcopy(tcga_metadataset)
          current_tcga_metadataset.task_ids = metadataset
          if preloaded:
              current_tcga_metadataset.open()
          tcga_metadatasets.append(current_tcga_metadataset)
  
      return tcga_metadatasets
  
imports: import os
  import json
  import h5py
  import numpy
  import torch
  import copy
  from ordered_set import OrderedSet
  from torchmeta.utils.data import Task, MetaDataset
  from torchmeta.datasets.utils import get_asset
label_desc: Use image node_id 3 for calling a global function with example usage: _split_tcga(tcga_metadataset, counts) and returns: tcga_metadatasets