class_docstr: EM algorithm used to quantize the columns of W to minimize
  
                       ||W - W_hat||^2
  
  Args:
      - W: weight matrix of size (in_features x out_features)
      - n_iter: number of k-means iterations
      - n_centroids: number of centroids (size of codebook)
      - eps: for cluster reassignment when an empty cluster is found
      - max_tentatives for cluster reassignment when an empty cluster is found
      - verbose: print error after each iteration
  
  Remarks:
      - If one cluster is empty, the most populated cluster is split into
        two clusters
      - All the relevant dimensions are specified in the code
docstr: There are two standard steps for each iteration: expectation (E) and
  minimization (M). The E-step (assignment) is performed with an exhaustive
  search and the M-step (centroid computation) is performed with
  the exact solution.
  
  Args:
      - i: step number
  
  Remarks:
      - The E-step heavily uses PyTorch broadcasting to speed up computations
        and reduce the memory overhead
docstr_tok: ["There","are","two","standard","steps","for","each","iteration",":","expectation","(","E",")","and","minimization","(","M",")",".","The","E-step","(","assignment",")","is","performed","with","an","exhaustive","search","and","the","M-step","(","centroid","computation",")","is","performed","with","the","exact","solution",".","Args",":","-","i",":","step","number","Remarks",":","-","The","E-step","heavily","uses","PyTorch","broadcasting","to","speed","up","computations","and","reduce","the","memory","overhead"]
code: def step(self, i):
      """
      There are two standard steps for each iteration: expectation (E) and
      minimization (M). The E-step (assignment) is performed with an exhaustive
      search and the M-step (centroid computation) is performed with
      the exact solution.
  
      Args:
          - i: step number
  
      Remarks:
          - The E-step heavily uses PyTorch broadcasting to speed up computations
            and reduce the memory overhead
      """
  
      # assignments (E-step)
      distances = (
          self.compute_distances()
      )  # (n_centroids x out_features)
      self.assignments = torch.argmin(
          distances, dim=0
      )  # (out_features)
      n_empty_clusters = self.resolve_empty_clusters()
  
      # centroids (M-step)
      for k in range(self.n_centroids):
          W_k = self.W[
              :, self.assignments == k
          ]  # (in_features x size_of_cluster_k)
          self.centroids[k] = W_k.mean(dim=1)  # (in_features)
  
      # book-keeping
      obj = (
          (self.centroids[self.assignments].t() - self.W)
          .norm(p=2)
          .item()
      )
      self.objective.append(obj)
      if self.verbose:
          logging.info(
              f"Iteration: {i},\t"
              f"objective: {obj:.6f},\t"
              f"resolved empty clusters: {n_empty_clusters}"
          )
  
file_path: fairseq/fairseq/modules/quantization/pq/em.py
imports: import logging
  import os
  import random
  from collections import Counter
  import torch
label_desc: Use image node_id 3 for calling the EM obj's underlying member method code with example usage: obj.step(i) without return types