code: def get_irrelevant_regions(
      labels, K, num_classes, N_shared=5, stage="TRAIN"
  ):
      """This class returns binary matrix that indicates the irrelevant regions in the VQ-dictionary given the labels array
  
      Arguments
      ---------
      labels : torch.tensor
          1 dimensional torch.tensor of size [B]
      K : int
          Number of keys in the dictionary
      num_classes : int
          Number of possible classes
      N_shared : int
          Number of shared keys
      stage : str
          "TRAIN" or else
  
      Example:
      --------
      >>> labels = torch.Tensor([1, 0, 2])
      >>> irrelevant_regions = get_irrelevant_regions(labels, 20, 3, 5)
      >>> print(irrelevant_regions.shape)
      torch.Size([3, 20])
      """
  
      uniform_mat = torch.round(
          torch.linspace(-0.5, num_classes - 0.51, K - N_shared)
      ).to(labels.device)
  
      uniform_mat = uniform_mat.unsqueeze(0).repeat(labels.shape[0], 1)
  
      labels_expanded = labels.unsqueeze(1).repeat(1, K - N_shared)
  
      irrelevant_regions = uniform_mat != labels_expanded
  
      if stage == "TRAIN":
          irrelevant_regions = (
              torch.cat(
                  [
                      irrelevant_regions,
                      torch.ones(
                          irrelevant_regions.shape[0], N_shared
                      ).to(labels.device),
                  ],
                  dim=1,
              )
              == 1
          )
      else:
          irrelevant_regions = (
              torch.cat(
                  [
                      irrelevant_regions,
                      torch.zeros(
                          irrelevant_regions.shape[0], N_shared
                      ).to(labels.device),
                  ],
                  dim=1,
              )
              == 1
          )
      return irrelevant_regions
  
imports: import torch
  import torch.nn
  from torch.autograd import Function
label_desc: Use image node_id 1 for calling a global function with example usage: get_irrelevant_regions(labels, K, num_classes, N_shared, stage) and returns: irrelevant_regions