code: def forward(self, input, return_soft_codes=False):
      if self.balancing_heuristic and self.codes_full:
          h = torch.histc(
              self.codes, bins=self.n_embed, min=0, max=self.n_embed
          ) / len(self.codes)
          mask = torch.logical_or(h > 0.9, h < 0.01).unsqueeze(1)
          ep = self.embed.permute(1, 0)
          ea = self.embed_avg.permute(1, 0)
          rand_embed = torch.randn_like(ep) * mask
          self.embed = (ep * ~mask + rand_embed).permute(1, 0)
          self.embed_avg = (ea * ~mask + rand_embed).permute(1, 0)
          self.cluster_size = self.cluster_size * ~mask.squeeze()
          if torch.any(mask):
              print(f"Reset {torch.sum(mask)} embedding codes.")
              self.codes = None
              self.codes_full = False
  
      flatten = input.reshape(-1, self.dim)
      dist = (
          flatten.pow(2).sum(1, keepdim=True)
          - 2 * flatten @ self.embed
          + self.embed.pow(2).sum(0, keepdim=True)
      )
      soft_codes = -dist
      _, embed_ind = soft_codes.max(1)
      embed_onehot = F.one_hot(embed_ind, self.n_embed).type(
          flatten.dtype
      )
      embed_ind = embed_ind.view(*input.shape[:-1])
      quantize = self.embed_code(embed_ind)
  
      if self.balancing_heuristic:
          if self.codes is None:
              self.codes = embed_ind.flatten()
          else:
              self.codes = torch.cat([self.codes, embed_ind.flatten()])
              if len(self.codes) > self.max_codes:
                  self.codes = self.codes[-self.max_codes :]
                  self.codes_full = True
  
      if self.training:
          embed_onehot_sum = embed_onehot.sum(0)
          embed_sum = flatten.transpose(0, 1) @ embed_onehot
  
          if (
              distributed.is_initialized()
              and distributed.get_world_size() > 1
          ):
              distributed.all_reduce(embed_onehot_sum)
              distributed.all_reduce(embed_sum)
  
          self.cluster_size.data.mul_(self.decay).add_(
              embed_onehot_sum, alpha=1 - self.decay
          )
          self.embed_avg.data.mul_(self.decay).add_(
              embed_sum, alpha=1 - self.decay
          )
          n = self.cluster_size.sum()
          cluster_size = (
              (self.cluster_size + self.eps)
              / (n + self.n_embed * self.eps)
              * n
          )
          embed_normalized = self.embed_avg / cluster_size.unsqueeze(0)
          self.embed.data.copy_(embed_normalized)
  
      diff = (quantize.detach() - input).pow(2).mean()
      quantize = input + (quantize - input).detach()
  
      if return_soft_codes:
          return (
              quantize,
              diff,
              embed_ind,
              soft_codes.view(input.shape[:-1] + (-1,)),
          )
      elif self.new_return_order:
          return quantize, embed_ind, diff
      else:
          return quantize, diff, embed_ind
  
imports: import functools
  from math import sqrt
  import torch
  import torch.distributed
  import torch.nn
  import torch.nn.functional
  import torchaudio
  from einops import rearrange
label_desc: Use image node_id 2 for calling the Quantize obj's underlying member method code with example usage: obj.forward(input, return_soft_codes) and returns: quantize,  diff,  embed_ind,  soft_codes, quantize,  embed_ind,  diff, quantize,  diff,  embed_ind