class_docstr: Codebook with Euclidean distance.
  
  Args:
      dim (int): Dimension.
      codebook_size (int): Codebook size.
      kmeans_init (bool): Whether to use k-means to initialize the codebooks.
          If set to true, run the k-means algorithm on the first training batch and use
          the learned centroids as initialization.
      kmeans_iters (int): Number of iterations used for k-means algorithm at initialization.
      decay (float): Decay for exponential moving average over the codebooks.
      epsilon (float): Epsilon value for numerical stability.
      threshold_ema_dead_code (int): Threshold for dead code expiration. Replace any codes
          that have an exponential moving average cluster size less than the specified threshold with
          randomly selected vector from the current batch.
class_docstr_tok: ["Codebook","with","Euclidean","distance",".","Args",":","dim","(","int",")",":","Dimension",".","codebook_size","(","int",")",":","Codebook","size",".","kmeans_init","(","bool",")",":","Whether","to","use","k-means","to","initialize","the","codebooks",".","If","set","to","true",",","run","the","k-means","algorithm","on","the","first","training","batch","and","use","the","learned","centroids","as","initialization",".","kmeans_iters","(","int",")",":","Number","of","iterations","used","for","k-means","algorithm","at","initialization",".","decay","(","float",")",":","Decay","for","exponential","moving","average","over","the","codebooks",".","epsilon","(","float",")",":","Epsilon","value","for","numerical","stability",".","threshold_ema_dead_code","(","int",")",":","Threshold","for","dead","code","expiration",".","Replace","any","codes","that","have","an","exponential","moving","average","cluster","size","less","than","the","specified","threshold","with","randomly","selected","vector","from","the","current","batch","."]
code: def replace_(self, samples, mask):
      modified_codebook = torch.where(
          mask[..., None],
          sample_vectors(samples, self.codebook_size),
          self.embed,
      )
      self.embed.data.copy_(modified_codebook)
  
imports: import typing
  from einops import rearrange, repeat
  import flashy
  import torch
  from torch import nn, einsum
  import torch.nn.functional
label_desc: Use image node_id 3 for calling the EuclideanCodebook obj's underlying member method code with example usage: obj.replace_(samples, mask) without return types