vector-quantize-pytorch 1.26.0__tar.gz → 1.27.0__tar.gz

{vector_quantize_pytorch-1.26.0 → vector_quantize_pytorch-1.27.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: vector-quantize-pytorch
-Version: 1.26.0
+Version: 1.27.0
 Summary: Vector Quantization - Pytorch
 Project-URL: Homepage, https://pypi.org/project/vector-quantize-pytorch/
 Project-URL: Repository, https://github.com/lucidrains/vector-quantizer-pytorch

{vector_quantize_pytorch-1.26.0 → vector_quantize_pytorch-1.27.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "vector-quantize-pytorch"
-version = "1.26.0"
+version = "1.27.0"
 description = "Vector Quantization - Pytorch"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }

vector_quantize_pytorch-1.27.0/tests/test_beam.py

@@ -0,0 +1,64 @@
+import torch
+from vector_quantize_pytorch import VectorQuantize
+
+def test_topk_and_manual_ema_update():
+
+    vq1 = VectorQuantize(
+        dim = 256,
+        codebook_size = 512
+    )
+
+    vq2 = VectorQuantize(
+        dim = 256,
+        codebook_size = 512
+    )
+    
+    vq2.load_state_dict(vq1.state_dict())
+
+    x = torch.randn(1, 1024, 256)
+    mask = torch.randint(0, 2, (1, 1024)).bool()
+
+    vq1.train()
+    quantize1, indices1, commit_loss1 = vq1(x, mask = mask)
+
+    vq2.train()
+    quantize2, indices2, commit_losses = vq2(x, mask = mask, topk = 1, ema_update = False)
+
+    assert quantize2.shape == (1, 1024, 1, 256)
+    assert indices2.shape == (1, 1024, 1)
+    assert commit_losses.shape == (1, 1024, 1)
+
+    top_quantize2 = quantize2[..., 0, :]
+    top_indices2 = indices2[..., 0]
+
+    assert torch.allclose(commit_loss1, commit_losses.sum() / mask.sum())
+    assert torch.equal(indices1, top_indices2)
+    assert torch.allclose(quantize1, top_quantize2)
+
+    assert not torch.allclose(vq1._codebook.embed_avg, vq2._codebook.embed_avg)
+
+    vq2.update_ema_indices(x, top_indices2, mask = mask)
+
+    assert torch.allclose(vq1._codebook.cluster_size, vq2._codebook.cluster_size)
+    assert torch.allclose(vq1._codebook.embed_avg, vq2._codebook.embed_avg)
+    assert torch.allclose(vq1.codebook, vq2.codebook)
+
+def test_beam_search():
+    import torch
+    from vector_quantize_pytorch import ResidualVQ
+
+    residual_vq = ResidualVQ(
+        dim = 256,
+        num_quantizers = 8,      # specify number of quantizers
+        codebook_size = 1024,    # codebook size
+        quantize_dropout = True
+    )
+
+    x = torch.randn(1, 1024, 256)
+
+    for _ in range(5):
+        quantized, indices, commit_loss = residual_vq(x, beam_size = 3)
+
+    assert quantized.shape == (1, 1024, 256)
+    assert indices.shape == (1, 1024, 8)
+    assert commit_loss.shape == (8,)

{vector_quantize_pytorch-1.26.0 → vector_quantize_pytorch-1.27.0}/vector_quantize_pytorch/residual_vq.py

@@ -6,7 +6,7 @@ from functools import partial, cache
 from itertools import zip_longest
 
 import torch
-from torch import nn, Tensor
+from torch import nn, Tensor, arange, cat
 from torch.nn import Module, ModuleList
 import torch.nn.functional as F
 import torch.distributed as dist
@@ -14,6 +14,7 @@ from vector_quantize_pytorch.vector_quantize_pytorch import VectorQuantize
 
 from einops import rearrange, repeat, reduce, pack, unpack
 
+import einx
 from einx import get_at
 
 # helper functions
@@ -36,6 +37,47 @@ def unique(arr):
 def round_up_multiple(num, mult):
     return ceil(num / mult) * mult
 
+# tensor helpers
+
+def pad_at_dim(
+    t,
+    pad: tuple[int, int],
+    dim = -1,
+    value = 0.
+):
+    if pad == (0, 0):
+        return t
+
+    dims_from_right = (- dim - 1) if dim < 0 else (t.ndim - dim - 1)
+    zeros = ((0, 0) * dims_from_right)
+    return F.pad(t, (*zeros, *pad), value = value)
+
+def pack_one(t, pattern):
+    packed, packed_shape = pack([t], pattern)
+
+    def inverse(out, inv_pattern = None):
+        inv_pattern = default(inv_pattern, pattern)
+        return first(unpack(out, packed_shape, inv_pattern))
+
+    return packed, inverse
+
+def batch_select(t, indices, pattern = None):
+
+    if exists(pattern):
+        indices = rearrange(indices, '... k -> (...) k')
+        t, inv_pack = pack_one(t, pattern)
+
+
+    batch_indices = arange(t.shape[0], device = t.device)
+    batch_indices = rearrange(batch_indices, 'b -> b 1')
+
+    out = t[batch_indices, indices]
+
+    if exists(pattern):
+        out = inv_pack(out)
+
+    return out
+
 # distributed helpers
 
 def is_distributed():
@@ -128,6 +170,8 @@ class ResidualVQ(Module):
         accept_image_fmap = False,
         implicit_neural_codebook = False, # QINCo from https://arxiv.org/abs/2401.14732
         mlp_kwargs: dict = dict(),
+        beam_size = None,
+        eval_beam_size = None,
         **vq_kwargs
     ):
         super().__init__()
@@ -183,6 +227,15 @@ class ResidualVQ(Module):
         self.quantize_dropout_cutoff_index = quantize_dropout_cutoff_index
         self.quantize_dropout_multiple_of = quantize_dropout_multiple_of  # encodec paper proposes structured dropout, believe this was set to 4
 
+        # determine whether is using ema update
+
+        self.vq_is_ema_updating = first(self.layers).ema_update
+
+        # beam size
+
+        self.beam_size = default(beam_size, eval_beam_size)
+        self.eval_beam_size = eval_beam_size
+
         # setting up the MLPs for implicit neural codebooks
         
         self.mlps = None
@@ -295,19 +348,29 @@ class ResidualVQ(Module):
         return_all_codes = False,
         sample_codebook_temp = None,
         freeze_codebook = False,
+        beam_size = None,
         rand_quantize_dropout_fixed_seed = None
     ):
-        num_quant, quant_dropout_multiple_of, return_loss, device = self.num_quantizers, self.quantize_dropout_multiple_of, exists(indices), x.device
+
+        # variables
+
+        input_shape, num_quant, quant_dropout_multiple_of, return_loss, device = x.shape, self.num_quantizers, self.quantize_dropout_multiple_of, exists(indices), x.device
+
+        beam_size = default(beam_size, self.beam_size if self.training else self.eval_beam_size)
+
+        is_beam_search = exists(beam_size) and beam_size > 1
+
+        # projecting in
 
         x = self.project_in(x)
 
         assert not (self.accept_image_fmap and exists(indices))
 
-        quantized_out = 0.
+        quantized_out = torch.zeros_like(x)
         residual = x
 
-        all_losses = []
-        all_indices = []
+        all_losses = torch.empty((0,), device = device, dtype = torch.float32)
+        all_indices = torch.empty((*input_shape[:-1], 0), device = device, dtype = torch.long)
 
         if isinstance(indices, list):
             indices = torch.stack(indices)
@@ -319,7 +382,6 @@ class ResidualVQ(Module):
         should_quantize_dropout = self.training and self.quantize_dropout and not return_loss
 
         # sample a layer index at which to dropout further residual quantization
-        # also prepare null indices and loss
 
         if should_quantize_dropout:
 
@@ -335,9 +397,24 @@ class ResidualVQ(Module):
             if quant_dropout_multiple_of != 1:
                 rand_quantize_dropout_index = round_up_multiple(rand_quantize_dropout_index + 1, quant_dropout_multiple_of) - 1
 
-            null_indices_shape = (x.shape[0], *x.shape[-2:]) if self.accept_image_fmap else tuple(x.shape[:2])
-            null_indices = torch.full(null_indices_shape, -1., device = device, dtype = torch.long)
-            null_loss = torch.full((1,), 0., device = device, dtype = x.dtype)
+        # save all inputs across layers, for use during expiration at end under shared codebook setting, or ema update during beam search
+
+        all_residuals = torch.empty((*input_shape[:-1], 0, input_shape[-1]), dtype = residual.dtype, device = device)
+
+        # maybe prepare beam search
+
+        if is_beam_search:
+            prec_dims = x.shape[:-1]
+
+            search_scores = torch.zeros((*prec_dims, 1), device = device, dtype = x.dtype)
+
+            residual = rearrange(residual, '... d -> ... 1 d')
+            quantized_out = rearrange(quantized_out, '... d -> ... 1 d')
+
+            all_residuals = rearrange(all_residuals, '... l d -> ... 1 l d')
+            all_indices = rearrange(all_indices, '... l -> ... 1 l')
+
+            all_losses = all_losses.reshape(*input_shape[:-1], 1, 0)
 
         # setup the mlps for implicit neural codebook
 
@@ -346,17 +423,13 @@ class ResidualVQ(Module):
         if self.implicit_neural_codebook:
             maybe_code_transforms = (None, *self.mlps)
 
-        # save all inputs across layers, for use during expiration at end under shared codebook setting
-
-        all_residuals = []
-
         # go through the layers
 
         for quantizer_index, (vq, maybe_mlp) in enumerate(zip(self.layers, maybe_code_transforms)):
 
             if should_quantize_dropout and quantizer_index > rand_quantize_dropout_index:
-                all_indices.append(null_indices)
-                all_losses.append(null_loss)
+                all_indices = pad_at_dim(all_indices, (0, 1), value = -1, dim = -1)
+                all_losses = pad_at_dim(all_losses, (0, 1), value = 0, dim = -1)
                 continue
 
             layer_indices = None
@@ -368,9 +441,9 @@ class ResidualVQ(Module):
             if exists(maybe_mlp):
                 maybe_mlp = partial(maybe_mlp, condition = quantized_out)
 
-            # save for expiration
+            # save the residual input for maybe expiration as well as ema update after beam search
 
-            all_residuals.append(residual)
+            all_residuals = cat((all_residuals, rearrange(residual, '... d -> ... 1 d')), dim = -2)
 
             # vector quantize forward
 
@@ -380,29 +453,111 @@ class ResidualVQ(Module):
                 indices = layer_indices,
                 sample_codebook_temp = sample_codebook_temp,
                 freeze_codebook = freeze_codebook,
-                codebook_transform_fn = maybe_mlp
+                codebook_transform_fn = maybe_mlp,
+                topk = beam_size if is_beam_search else None
             )
 
-            residual = residual - quantized.detach()
-            quantized_out = quantized_out + quantized
+            # cross entropy loss for some old paper
 
             if return_loss:
-                ce_loss = rest[0]
+                ce_loss = first(rest)
                 ce_losses.append(ce_loss)
                 continue
 
             embed_indices, loss = rest
 
-            all_indices.append(embed_indices)
-            all_losses.append(loss)
+            # handle expanding first residual if doing beam search
+
+            if is_beam_search:
+
+                search_scores = einx.add('... j, ... j k -> ... (j k)', search_scores, -loss)
+
+                residual = rearrange(residual, '... j d -> ... j 1 d')
+                quantized_out = rearrange(quantized_out, '... j d -> ... j 1 d')
+
+                all_residuals = repeat(all_residuals, '... j l d -> ... (j k) l d', k = beam_size)
+
+            # core residual vq logic
+
+            residual = residual - quantized.detach()
+            quantized_out = quantized_out + quantized
+
+            # handle sort and topk beams
+
+            if is_beam_search:
+                residual = rearrange(residual, '... j k d -> ... (j k) d')
+                quantized_out = rearrange(quantized_out, '... j k d -> ... (j k) d')
+
+                # broadcat the indices
+
+                all_indices = repeat(all_indices, '... j l -> ... j k l', k = embed_indices.shape[-1])
+                embed_indices = rearrange(embed_indices, '... j k -> ... j k 1')
+
+                all_indices = cat((all_indices, embed_indices), dim = -1)
+                all_indices = rearrange(all_indices, '... j k l -> ... (j k) l')
+
+                # broadcat the losses
+
+                all_losses = repeat(all_losses, '... j l -> ... j k l', k = loss.shape[-1])
+                loss = rearrange(loss, '... -> ... 1')
+
+                all_losses = cat((all_losses, loss), dim = -1)
+                all_losses = rearrange(all_losses, '... j k l -> ... (j k) l')
+
+                # handle sort and selection of highest beam size
+
+                if search_scores.shape[-1] > beam_size:
+                    search_scores, select_indices = search_scores.topk(beam_size, dim = -1)
+
+                    residual = batch_select(residual, select_indices, '* k d')
+                    quantized_out = batch_select(quantized_out, select_indices, '* k d')
+
+                    all_indices = batch_select(all_indices, select_indices, '* k l')
+                    all_losses = batch_select(all_losses, select_indices, '* k l')
+
+                    all_residuals = batch_select(all_residuals, select_indices, '* k l d')
+            else:
+                # aggregate indices and losses
+
+                all_indices = cat((all_indices, rearrange(embed_indices, '... -> ... 1')), dim = -1)
+
+                all_losses = cat((all_losses, rearrange(loss, '... -> ... 1')), dim = -1)
+
+        # handle beam search
+
+        if is_beam_search:
+            top_index = search_scores.argmax(dim = -1, keepdim = True)
+
+            quantized_out = batch_select(quantized_out, top_index, '* k d')
+            all_indices = batch_select(all_indices, top_index, '* k l')
+            all_losses = batch_select(all_losses, top_index, '* k l')
+            all_residuals = batch_select(all_residuals, top_index, '* k l d')
+
+            quantized_out, all_indices, all_losses, all_residuals = [t[..., 0, :] for t in (quantized_out, all_indices, all_losses, all_residuals)]
+
+            # handle commit loss, which should be the average
+
+            if exists(mask):
+                all_losses = einx.where('..., ... l,', mask, all_losses, 0.)
+                all_losses = reduce(all_losses, '... l -> l', 'sum') / mask.sum(dim = -1).clamp_min(1e-4)
+            else:
+                all_losses = reduce(all_losses, '... l -> l', 'mean')
+
+            # handle updating ema
+
+            if self.vq_is_ema_updating:
+                for vq, layer_input, indices in zip(self.layers, all_residuals.unbind(dim = -2), all_indices.unbind(dim = -1)): # in the case of quantize dropout, zip will terminate with the shorter sequence, which should be all_residuals
+                    vq.update_ema_indices(layer_input, indices, mask = mask)
 
         # if shared codebook, update ema only at end
 
-        if self.training and self.shared_codebook:
+        if self.training and self.shared_codebook and not is_beam_search:
             shared_layer = first(self.layers)
             shared_layer._codebook.update_ema()
             shared_layer.update_in_place_optimizer()
-            shared_layer.expire_codes_(torch.cat(all_residuals, dim = -2))
+
+            all_codes_for_expire = rearrange(all_residuals, '... n l d -> ... (n l) d')
+            shared_layer.expire_codes_(all_codes_for_expire)
 
         # project out, if needed
 
@@ -415,8 +570,6 @@ class ResidualVQ(Module):
 
         # stack all losses and indices
 
-        all_losses, all_indices = map(partial(torch.stack, dim = -1), (all_losses, all_indices))
-
         ret = (quantized_out, all_indices, all_losses)
 
         if return_all_codes:

{vector_quantize_pytorch-1.26.0 → vector_quantize_pytorch-1.27.0}/vector_quantize_pytorch/vector_quantize_pytorch.py

@@ -1,4 +1,5 @@
 from __future__ import annotations
+from typing import Callable
 
 from math import sqrt
 from functools import partial, cache
@@ -6,7 +7,7 @@ from collections import namedtuple
 
 import torch
 from torch.nn import Module
-from torch import nn, einsum, is_tensor, Tensor
+from torch import nn, einsum, tensor, is_tensor, Tensor
 import torch.nn.functional as F
 import torch.distributed as distributed
 from torch.optim import Optimizer
@@ -15,8 +16,6 @@ from torch.amp import autocast
 import einx
 from einops import rearrange, repeat, reduce, pack, unpack
 
-from typing import Callable
-
 def exists(val):
     return val is not None
 
@@ -120,7 +119,8 @@ def gumbel_sample(
     stochastic = False,
     straight_through = False,
     dim = -1,
-    training = True
+    training = True,
+    topk = None
 ):
     dtype, size = logits.dtype, logits.shape[dim]
 
@@ -129,7 +129,11 @@ def gumbel_sample(
     else:
         sampling_logits = logits
 
-    ind = sampling_logits.argmax(dim = dim)
+    if exists(topk):
+        ind = sampling_logits.topk(topk, dim = dim).indices
+    else:
+        ind = sampling_logits.argmax(dim = dim)
+
     one_hot = F.one_hot(ind, size).type(dtype)
 
     if not straight_through or temperature <= 0. or not training:
@@ -182,7 +186,7 @@ def sample_multinomial(total_count, probs):
     return sample.to(device)
 
 def all_gather_sizes(x, dim):
-    size = torch.tensor(x.shape[dim], dtype = torch.long, device = x.device)
+    size = tensor(x.shape[dim], dtype = torch.long, device = x.device)
     all_sizes = [torch.empty_like(size) for _ in range(distributed.get_world_size())]
     distributed.all_gather(all_sizes, size)
     return torch.stack(all_sizes)
@@ -337,7 +341,7 @@ def orthogonal_loss_fn(t):
 
 # distance types
 
-class EuclideanCodebook(Module):
+class Codebook(Module):
     def __init__(
         self,
         dim,
@@ -359,17 +363,23 @@ class EuclideanCodebook(Module):
         affine_param = False,
         sync_affine_param = False,
         affine_param_batch_decay = 0.99,
-        affine_param_codebook_decay = 0.9
+        affine_param_codebook_decay = 0.9,
+        use_cosine_sim = False
     ):
         super().__init__()
-        self.transform_input = identity
+        self.transform_input = identity if not use_cosine_sim else l2norm
 
         self.decay = decay
         self.ema_update = ema_update
         self.manual_ema_update = manual_ema_update
 
-        init_fn = uniform_init if not kmeans_init else torch.zeros
-        embed = init_fn(num_codebooks, codebook_size, dim)
+        if kmeans_init:
+            embed = torch.zeros(num_codebooks, codebook_size, dim)
+        else:
+            embed = uniform_init(num_codebooks, codebook_size, dim)
+
+            if use_cosine_sim:
+                embed = l2norm(embed)
 
         self.codebook_size = codebook_size
         self.num_codebooks = num_codebooks
@@ -392,7 +402,7 @@ class EuclideanCodebook(Module):
         self.kmeans_all_reduce_fn = distributed.all_reduce if use_ddp and sync_kmeans else noop
         self.all_reduce_fn = distributed.all_reduce if use_ddp else noop
 
-        self.register_buffer('initted', torch.Tensor([not kmeans_init]))
+        self.register_buffer('initted', tensor(not kmeans_init))
         self.register_buffer('cluster_size', torch.ones(num_codebooks, codebook_size))
         self.register_buffer('embed_avg', embed.clone())
 
@@ -402,6 +412,8 @@ class EuclideanCodebook(Module):
         else:
             self.register_buffer('embed', embed)
 
+        self.use_cosine_sim = use_cosine_sim
+
         # affine related params
 
         self.affine_param = affine_param
@@ -416,9 +428,9 @@ class EuclideanCodebook(Module):
         self.register_buffer('batch_mean', None)
         self.register_buffer('batch_variance', None)
 
-        self.register_buffer('codebook_mean_needs_init', torch.Tensor([True]))
+        self.register_buffer('codebook_mean_needs_init', tensor(True))
         self.register_buffer('codebook_mean', torch.empty(num_codebooks, 1, dim))
-        self.register_buffer('codebook_variance_needs_init', torch.Tensor([True]))
+        self.register_buffer('codebook_variance_needs_init', tensor(True))
         self.register_buffer('codebook_variance', torch.empty(num_codebooks, 1, dim))
 
     @torch.jit.ignore
@@ -434,6 +446,7 @@ class EuclideanCodebook(Module):
             data,
             self.codebook_size,
             self.kmeans_iters,
+            use_cosine_sim = self.use_cosine_sim,
             sample_fn = self.sample_fn,
             all_reduce_fn = self.kmeans_all_reduce_fn
         )
@@ -443,7 +456,7 @@ class EuclideanCodebook(Module):
         self.embed_avg.data.copy_(embed_sum)
         self.cluster_size.data.copy_(cluster_size)
         self.update_ema()
-        self.initted.data.copy_(torch.Tensor([True]))
+        self.initted.data.copy_(tensor(True))
 
     @torch.jit.ignore
     def update_with_decay(self, buffer_name, new_value, decay):
@@ -452,7 +465,7 @@ class EuclideanCodebook(Module):
         needs_init = getattr(self, buffer_name + "_needs_init", False)
 
         if needs_init:
-            self.register_buffer(buffer_name + "_needs_init", torch.Tensor([False]))
+            self.register_buffer(buffer_name + "_needs_init", tensor(False))
 
         if not exists(old_value) or needs_init:
             self.register_buffer(buffer_name, new_value.detach())
@@ -495,7 +508,7 @@ class EuclideanCodebook(Module):
 
         # number of vectors, for denominator
 
-        num_vectors = torch.tensor([num_vectors], device = device, dtype = dtype)
+        num_vectors = tensor(num_vectors, device = device, dtype = dtype)
         distributed.all_reduce(num_vectors)
 
         # calculate distributed mean
@@ -515,6 +528,9 @@ class EuclideanCodebook(Module):
         self.update_with_decay('batch_variance', batch_variance, self.affine_param_batch_decay)
 
     def replace(self, batch_samples, batch_mask):
+        if self.use_cosine_sim:
+            batch_samples = l2norm(batch_samples)
+
         for ind, (samples, mask) in enumerate(zip(batch_samples, batch_mask)):
             sampled = self.replace_sample_fn(rearrange(samples, '... -> 1 ...'), mask.sum().item())
             sampled = rearrange(sampled, '1 ... -> ...')
@@ -539,247 +555,74 @@ class EuclideanCodebook(Module):
         cluster_size = laplace_smoothing(self.cluster_size, self.codebook_size, self.eps) * self.cluster_size.sum(dim = -1, keepdim = True)
 
         embed_normalized = self.embed_avg / rearrange(cluster_size, '... -> ... 1')
+
+        if self.use_cosine_sim:
+            embed_normalized = l2norm(embed_normalized)
+
         self.embed.data.copy_(embed_normalized)
 
-    @autocast('cuda', enabled = False)
-    def forward(
+    def update_ema_part(
         self,
-        x,
-        sample_codebook_temp = None,
+        flatten,
+        embed_onehot,
         mask = None,
-        freeze_codebook = False,
-        codebook_transform_fn: Callable | None = None,
         ema_update_weight: Tensor | Callable | None = None,
         accum_ema_update = False
     ):
-        needs_codebook_dim = x.ndim < 4
-        sample_codebook_temp = default(sample_codebook_temp, self.sample_codebook_temp)
-
-        x = x.float()
-
-        if needs_codebook_dim:
-            x = rearrange(x, '... -> 1 ...')
-
-        dtype = x.dtype
-        flatten, unpack_one = pack_one(x, 'h * d')
-
-        if exists(mask):
-            mask = repeat(mask, 'b n -> c (b h n)', c = flatten.shape[0], h = flatten.shape[-2] // (mask.shape[0] * mask.shape[1]))
-
-        self.init_embed_(flatten, mask = mask)
-
-        if self.affine_param:
-            self.update_affine(flatten, self.embed, mask = mask)
-
-        # get maybe learnable codes
-
-        embed = self.embed if self.learnable_codebook else self.embed.detach()
-
-        embed = embed.to(dtype)
-
-        # affine params
         if self.affine_param:
             codebook_std = self.codebook_variance.clamp(min = 1e-5).sqrt()
             batch_std = self.batch_variance.clamp(min = 1e-5).sqrt()
-            embed = (embed - self.codebook_mean) * (batch_std / codebook_std) + self.batch_mean
-
-        # handle maybe implicit neural codebook
-        # and calculate distance
+            flatten = (flatten - self.batch_mean) * (codebook_std / batch_std) + self.codebook_mean
 
-        if exists(codebook_transform_fn):
-            transformed_embed = codebook_transform_fn(embed)
-            transformed_embed = rearrange(transformed_embed, 'h b n c d -> h (b n) c d')
-            broadcastable_input = rearrange(flatten, '... d -> ... 1 d')
-
-            dist = -F.pairwise_distance(broadcastable_input, transformed_embed)
-        else:
-            dist = -cdist(flatten, embed)
-
-        # sample or argmax depending on temperature
-
-        embed_ind, embed_onehot = self.gumbel_sample(dist, dim = -1, temperature = sample_codebook_temp, training = self.training)
-
-        embed_ind = unpack_one(embed_ind, 'h *')
+        if exists(mask):
+            embed_onehot[~mask] = 0.
 
-        if exists(codebook_transform_fn):
-            transformed_embed = unpack_one(transformed_embed, 'h * c d')
+        cluster_size = embed_onehot.sum(dim = 1)
+        self.all_reduce_fn(cluster_size)
 
-        if self.training:
-            unpacked_onehot = unpack_one(embed_onehot, 'h * c')
+        embed_sum = einsum('h n d, h n c -> h c d', flatten, embed_onehot)
+        embed_sum = embed_sum.contiguous()
+        self.all_reduce_fn(embed_sum)
 
-            if exists(codebook_transform_fn):
-                quantize = einsum('h b n c, h b n c d -> h b n d', unpacked_onehot, transformed_embed)
-            else:
-                quantize = einsum('h b n c, h c d -> h b n d', unpacked_onehot, embed)
+        if callable(ema_update_weight):
+            ema_update_weight = ema_update_weight(embed_sum, cluster_size)
 
+        if accum_ema_update:
+            accum_grad_(self.cluster_size, cluster_size)
+            accum_grad_(self.embed_avg, embed_sum)
         else:
-            if exists(codebook_transform_fn):
-                # quantize = einx.get_at('h b n [c] d, h b n -> h b n d', transformed_embed, embed_ind)
-
-                repeated_embed_ind = repeat(embed_ind, 'h b n -> h b n 1 d', d = transformed_embed.shape[-1])
-                quantize = transformed_embed.gather(-2, repeated_embed_ind)
-                quantize = rearrange(quantize, 'h b n 1 d -> h b n d')
-
-            else:
-                # quantize = einx.get_at('h [c] d, h b n -> h b n d', embed, embed_ind)
-
-                repeated_embed = repeat(embed, 'h c d -> h b c d', b = embed_ind.shape[1])
-                repeated_embed_ind = repeat(embed_ind, 'h b n -> h b n d', d = embed.shape[-1])
-                quantize = repeated_embed.gather(-2, repeated_embed_ind)
-
-        if self.training and self.ema_update and not freeze_codebook:
-
-            if self.affine_param:
-                flatten = (flatten - self.batch_mean) * (codebook_std / batch_std) + self.codebook_mean
-
-            if exists(mask):
-                embed_onehot[~mask] = 0.
-
-            cluster_size = embed_onehot.sum(dim = 1)
-            self.all_reduce_fn(cluster_size)
+            ema_inplace(self.cluster_size, cluster_size, self.decay, ema_update_weight)
+            ema_inplace(self.embed_avg, embed_sum, self.decay, ema_update_weight)
 
-            embed_sum = einsum('h n d, h n c -> h c d', flatten, embed_onehot)
-            embed_sum = embed_sum.contiguous()
-            self.all_reduce_fn(embed_sum)
+            if not self.manual_ema_update:
+                self.update_ema()
+                self.expire_codes_(flatten)
 
-            if callable(ema_update_weight):
-                ema_update_weight = ema_update_weight(embed_sum, cluster_size)
-
-            if accum_ema_update:
-                accum_grad_(self.cluster_size, cluster_size)
-                accum_grad_(self.embed_avg, embed_sum)
-            else:
-                ema_inplace(self.cluster_size, cluster_size, self.decay, ema_update_weight)
-                ema_inplace(self.embed_avg, embed_sum, self.decay, ema_update_weight)
-
-                if not self.manual_ema_update:
-                    self.update_ema()
-                    self.expire_codes_(x)
-
-        if needs_codebook_dim:
-            quantize, embed_ind = map(lambda t: rearrange(t, '1 ... -> ...'), (quantize, embed_ind))
-
-        dist = unpack_one(dist, 'h * d')
-
-        return quantize, embed_ind, dist
-
-class CosineSimCodebook(Module):
-    def __init__(
+    def update_ema_indices(
         self,
-        dim,
-        codebook_size,
-        num_codebooks = 1,
-        kmeans_init = False,
-        kmeans_iters = 10,
-        sync_kmeans = True,
-        decay = 0.8,
-        eps = 1e-5,
-        threshold_ema_dead_code = 2,
-        reset_cluster_size = None,
-        use_ddp = False,
-        learnable_codebook = False,
-        gumbel_sample = gumbel_sample,
-        sample_codebook_temp = 1.,
-        ema_update = True,
-        manual_ema_update = False
+        x,
+        embed_ind,
+        mask = None,
+        ema_update_weight: Tensor | Callable | None = None,
+        accum_ema_update = False
     ):
-        super().__init__()
-        self.transform_input = l2norm
-
-        self.ema_update = ema_update
-        self.manual_ema_update = manual_ema_update
-
-        self.decay = decay
-
-        if not kmeans_init:
-            embed = l2norm(uniform_init(num_codebooks, codebook_size, dim))
-        else:
-            embed = torch.zeros(num_codebooks, codebook_size, dim)
-
-        self.codebook_size = codebook_size
-        self.num_codebooks = num_codebooks
-
-        self.kmeans_iters = kmeans_iters
-        self.eps = eps
-        self.threshold_ema_dead_code = threshold_ema_dead_code
-        self.reset_cluster_size = default(reset_cluster_size, threshold_ema_dead_code)
-
-        assert callable(gumbel_sample)
-        self.gumbel_sample = gumbel_sample
-        self.sample_codebook_temp = sample_codebook_temp
-
-        self.sample_fn = sample_vectors_distributed if use_ddp and sync_kmeans else batched_sample_vectors
-
-        self.replace_sample_fn = sample_vectors_distributed if use_ddp and sync_kmeans else batched_sample_vectors
-
-        self.kmeans_all_reduce_fn = distributed.all_reduce if use_ddp and sync_kmeans else noop
-        self.all_reduce_fn = distributed.all_reduce if use_ddp else noop
-
-        self.register_buffer('initted', torch.Tensor([not kmeans_init]))
-        self.register_buffer('cluster_size', torch.ones(num_codebooks, codebook_size))
-        self.register_buffer('embed_avg', embed.clone())
+        needs_codebook_dim = x.ndim < 4
+        x = x.float()
 
-        self.learnable_codebook = learnable_codebook
-        if learnable_codebook:
-            self.embed = nn.Parameter(embed)
-        else:
-            self.register_buffer('embed', embed)
+        if needs_codebook_dim:
+            x = rearrange(x, '... -> 1 ...')
 
-    @torch.jit.ignore
-    def init_embed_(self, data, mask = None):
-        if self.initted:
-            return
+        dtype = x.dtype
+        flatten, unpack_one = pack_one(x, 'h * d')
 
         if exists(mask):
-            c = data.shape[0]
-            data = rearrange(data[mask], '(c n) d -> c n d', c = c)
-
-        embed, cluster_size = kmeans(
-            data,
-            self.codebook_size,
-            self.kmeans_iters,
-            use_cosine_sim = True,
-            sample_fn = self.sample_fn,
-            all_reduce_fn = self.kmeans_all_reduce_fn
-        )
-
-        embed_sum = embed * rearrange(cluster_size, '... -> ... 1')
-
-        self.embed_avg.data.copy_(embed_sum)
-        self.cluster_size.data.copy_(cluster_size)
-        self.update_ema()
-        self.initted.data.copy_(torch.Tensor([True]))
-
-    def replace(self, batch_samples, batch_mask):
-        batch_samples = l2norm(batch_samples)
-
-        for ind, (samples, mask) in enumerate(zip(batch_samples, batch_mask)):
-            sampled = self.replace_sample_fn(rearrange(samples, '... -> 1 ...'), mask.sum().item())
-            sampled = rearrange(sampled, '1 ... -> ...')
-
-            self.embed.data[ind][mask] = sampled
-            self.embed_avg.data[ind][mask] = sampled * self.reset_cluster_size
-            self.cluster_size.data[ind][mask] = self.reset_cluster_size
-
-    def expire_codes_(self, batch_samples):
-        if self.threshold_ema_dead_code == 0:
-            return
-
-        expired_codes = self.cluster_size < self.threshold_ema_dead_code
-
-        if not torch.any(expired_codes):
-            return
-
-        batch_samples = rearrange(batch_samples, 'h ... d -> h (...) d')
-        self.replace(batch_samples, batch_mask = expired_codes)
+            mask = repeat(mask, 'b n -> c (b h n)', c = flatten.shape[0], h = flatten.shape[-2] // (mask.shape[0] * mask.shape[1]))
 
-    def update_ema(self):
-        cluster_size = laplace_smoothing(self.cluster_size, self.codebook_size, self.eps) * self.cluster_size.sum(dim = -1, keepdim = True)
+        embed_ind, _ = pack([embed_ind], 'h *')
+        embed_ind = embed_ind.masked_fill(embed_ind == -1, 0)
+        embed_onehot = F.one_hot(embed_ind, self.codebook_size).type(dtype)
 
-        embed_normalized = self.embed_avg / rearrange(cluster_size, '... -> ... 1')
-        embed_normalized = l2norm(embed_normalized)
-
-        self.embed.data.copy_(embed_normalized)
+        self.update_ema_part(flatten, embed_onehot, mask = mask, ema_update_weight = ema_update_weight, accum_ema_update = accum_ema_update)
 
     @autocast('cuda', enabled = False)
     def forward(
@@ -789,9 +632,13 @@ class CosineSimCodebook(Module):
         mask = None,
         freeze_codebook = False,
         codebook_transform_fn: Callable | None = None,
-        ema_update_weight: Tensor | None = None,
-        accum_ema_update = False
+        ema_update_weight: Tensor | Callable | None = None,
+        accum_ema_update = False,
+        ema_update = None,
+        topk = None
     ):
+        ema_update = default(ema_update, self.ema_update)
+
         needs_codebook_dim = x.ndim < 4
         sample_codebook_temp = default(sample_codebook_temp, self.sample_codebook_temp)
 
@@ -801,7 +648,6 @@ class CosineSimCodebook(Module):
             x = rearrange(x, '... -> 1 ...')
 
         dtype = x.dtype
-
         flatten, unpack_one = pack_one(x, 'h * d')
 
         if exists(mask):
@@ -809,35 +655,63 @@ class CosineSimCodebook(Module):
 
         self.init_embed_(flatten, mask = mask)
 
+        if self.affine_param:
+            self.update_affine(flatten, self.embed, mask = mask)
+
+        # get maybe learnable codes
+
         embed = self.embed if self.learnable_codebook else self.embed.detach()
 
         embed = embed.to(dtype)
 
+        # affine params
+
+        if self.affine_param:
+            codebook_std = self.codebook_variance.clamp(min = 1e-5).sqrt()
+            batch_std = self.batch_variance.clamp(min = 1e-5).sqrt()
+            embed = (embed - self.codebook_mean) * (batch_std / codebook_std) + self.batch_mean
+
         # handle maybe implicit neural codebook
-        # and compute cosine sim distance
+        # and calculate distance
 
         if exists(codebook_transform_fn):
             transformed_embed = codebook_transform_fn(embed)
             transformed_embed = rearrange(transformed_embed, 'h b n c d -> h (b n) c d')
-            transformed_embed = l2norm(transformed_embed)
 
-            dist = einsum('h n d, h n c d -> h n c', flatten, transformed_embed)
+            if self.use_cosine_sim:
+                transformed_embed = l2norm(transformed_embed)
+                dist = einsum('h n d, h n c d -> h n c', flatten, transformed_embed)
+            else:
+                broadcastable_input = rearrange(flatten, '... d -> ... 1 d')
+                dist = -F.pairwise_distance(broadcastable_input, transformed_embed)
         else:
-            dist = einsum('h n d, h c d -> h n c', flatten, embed)
+            if self.use_cosine_sim:
+                 dist = einsum('h n d, h c d -> h n c', flatten, embed)
+            else:
+                 dist = -cdist(flatten, embed)
+
+        # sample or argmax depending on temperature
 
-        embed_ind, embed_onehot = self.gumbel_sample(dist, dim = -1, temperature = sample_codebook_temp, training = self.training)
-        embed_ind = unpack_one(embed_ind, 'h *')
+        embed_ind, embed_onehot = self.gumbel_sample(dist, dim = -1, topk = topk, temperature = sample_codebook_temp, training = self.training)
+
+        if exists(topk):
+            embed_ind = unpack_one(embed_ind, 'h * k')
+        else:
+            embed_ind = unpack_one(embed_ind, 'h *')
 
         if exists(codebook_transform_fn):
             transformed_embed = unpack_one(transformed_embed, 'h * c d')
 
         if self.training:
-            unpacked_onehot = unpack_one(embed_onehot, 'h * c')
+            if exists(topk):
+                unpacked_onehot = unpack_one(embed_onehot, 'h * k c')
+            else:
+                unpacked_onehot = unpack_one(embed_onehot, 'h * c')
 
             if exists(codebook_transform_fn):
-                quantize = einsum('h b n c, h b n c d -> h b n d', unpacked_onehot, transformed_embed)
+                quantize = einsum('h b n ... c, h b n c d -> h b n ... d', unpacked_onehot, transformed_embed)
             else:
-                quantize = einsum('h b n c, h c d -> h b n d', unpacked_onehot, embed)
+                quantize = einsum('h b n ... c, h c d -> h b n ... d', unpacked_onehot, embed)
 
         else:
             if exists(codebook_transform_fn):
@@ -854,36 +728,14 @@ class CosineSimCodebook(Module):
                 repeated_embed_ind = repeat(embed_ind, 'h b n -> h b n d', d = embed.shape[-1])
                 quantize = repeated_embed.gather(-2, repeated_embed_ind)
 
-        if self.training and self.ema_update and not freeze_codebook:
-            if exists(mask):
-                embed_onehot[~mask] = 0.
-
-            bins = embed_onehot.sum(dim = 1)
-            self.all_reduce_fn(bins)
-
-            embed_sum = einsum('h n d, h n c -> h c d', flatten, embed_onehot)
-            embed_sum = embed_sum.contiguous()
-            self.all_reduce_fn(embed_sum)
-
-            if callable(ema_update_weight):
-                ema_update_weight = ema_update_weight(embed_sum, bins)
-
-            if accum_ema_update:
-                accum_grad_(self.cluster_size, bins)
-                accum_grad_(self.embed_avg, embed_sum)
-            else:
-
-                ema_inplace(self.cluster_size, bins, self.decay, ema_update_weight)
-                ema_inplace(self.embed_avg, embed_sum, self.decay, ema_update_weight)
-
-                if not self.manual_ema_update:
-                    self.update_ema()
-                    self.expire_codes_(x)
+        if self.training and ema_update and not freeze_codebook and not exists(topk):
+            self.update_ema_part(flatten, embed_onehot, mask = mask, ema_update_weight = ema_update_weight, accum_ema_update = accum_ema_update)
 
         if needs_codebook_dim:
             quantize, embed_ind = map(lambda t: rearrange(t, '1 ... -> ...'), (quantize, embed_ind))
 
         dist = unpack_one(dist, 'h * d')
+
         return quantize, embed_ind, dist
 
 # main class
@@ -1001,7 +853,7 @@ class VectorQuantize(Module):
 
         self.sync_update_v = sync_update_v
 
-        codebook_class = EuclideanCodebook if not use_cosine_sim else CosineSimCodebook
+        codebook_class = Codebook
 
         gumbel_sample_fn = partial(
             gumbel_sample,
@@ -1027,7 +879,8 @@ class VectorQuantize(Module):
             sample_codebook_temp = sample_codebook_temp,
             gumbel_sample = gumbel_sample_fn,
             ema_update = ema_update,
-            manual_ema_update = manual_ema_update
+            manual_ema_update = manual_ema_update,
+            use_cosine_sim = use_cosine_sim
         )
 
         if affine_param:
@@ -1051,7 +904,7 @@ class VectorQuantize(Module):
         self.accept_image_fmap = accept_image_fmap
         self.channel_last = channel_last
 
-        self.register_buffer('zero', torch.tensor(0.), persistent = False)
+        self.register_buffer('zero', tensor(0.), persistent = False)
 
         # for variable lengthed sequences, whether to take care of masking out the padding to 0 (or return the original input)
         self.return_zeros_for_masked_padding = return_zeros_for_masked_padding
@@ -1059,6 +912,10 @@ class VectorQuantize(Module):
         # whether to freeze the codebook, can be overridden on forward
         self.freeze_codebook = freeze_codebook
 
+    @property
+    def ema_update(self):
+        return self._codebook.ema_update
+
     @property
     def codebook(self):
         codebook = self._codebook.embed
@@ -1133,18 +990,47 @@ class VectorQuantize(Module):
         x = self.maybe_split_heads_from_input(x)
         self._codebook.expire_codes_(x)
 
+    def update_ema_indices(self, x, indices, mask = None):
+        if self.accept_image_fmap:
+            assert not exists(mask)
+            height, width = x.shape[-2:]
+            x = rearrange(x, 'b c h w -> b (h w) c')
+
+        if not self.channel_last and not self.accept_image_fmap:
+            x = rearrange(x, 'b d n -> b n d')
+
+        x = self.project_in(x)
+        x = self.maybe_split_heads_from_input(x)
+        x = self._codebook.transform_input(x)
+
+        if self.heads > 1:
+            if self.separate_codebook_per_head:
+                indices = rearrange(indices, 'b n h -> h b n')
+            else:
+                indices = rearrange(indices, 'b n h -> 1 (b h) n')
+
+        if self.accept_image_fmap:
+             indices = rearrange(indices, 'b h w ... -> b (h w) ...')
+
+        if x.ndim == 2: # only one token
+             indices = rearrange(indices, 'b ... -> b 1 ...')
+
+        self._codebook.update_ema_indices(x, indices, mask = mask)
+
     def forward(
         self,
         x,
         indices = None,
         mask = None,
         lens = None,
+        topk = None,
         sample_codebook_temp = None,
         freeze_codebook = None,
         return_loss_breakdown = False,
         codebook_transform_fn: Callable | None = None,
         ema_update_weight: Tensor | None = None,
-        accum_ema_update = False
+        accum_ema_update = False,
+        ema_update = None
     ):
         orig_input, input_requires_grad = x, x.requires_grad
 
@@ -1202,7 +1088,9 @@ class VectorQuantize(Module):
             freeze_codebook = freeze_codebook,
             codebook_transform_fn = codebook_transform_fn,
             ema_update_weight = ema_update_weight,
-            accum_ema_update = accum_ema_update
+            accum_ema_update = accum_ema_update,
+            ema_update = ema_update and not exists(topk),
+            topk = topk
         )
 
         # quantize
@@ -1301,7 +1189,7 @@ class VectorQuantize(Module):
 
         # aggregate loss
 
-        loss = torch.tensor([0.], device = device, requires_grad = self.training)
+        loss = tensor(0., device = device, requires_grad = self.training)
 
         if self.training:
             # calculate codebook diversity loss (negative of entropy) if needed
@@ -1326,9 +1214,19 @@ class VectorQuantize(Module):
 
                     commit_loss = calculate_ce_loss(embed_ind)
                 else:
-                    if exists(mask):
+                    if exists(topk):
+                        # handle special case when returning topk
+
+                        repeated_input = repeat(orig_input, '... d -> ... k d', k = topk)
+                        commit_loss = F.mse_loss(commit_quantize, repeated_input, reduction = 'none')
+                        commit_loss = reduce(commit_loss, '... k d -> ... k', 'mean')
+
+                        if exists(mask):
+                            commit_loss = einx.where('..., ... k, -> ... k', mask, commit_loss, 0.)
+
+                    elif exists(mask):
                         # with variable lengthed sequences
-                        commit_loss = F.mse_loss(commit_quantize, x, reduction = 'none')
+                        commit_loss = F.mse_loss(commit_quantize, orig_input, reduction = 'none')
 
                         loss_mask = mask
                         if is_multiheaded:
@@ -1391,7 +1289,7 @@ class VectorQuantize(Module):
                 masked_out_value = torch.zeros_like(orig_input)
 
             quantize = einx.where(
-                'b n, b n d, b n d -> b n d',
+                'b n, b n ... d, b n d -> b n ... d',
                 mask,
                 quantize,
                 masked_out_value