revert uneeded fix #120

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "vector-quantize-pytorch"
-version = "1.14.11"
+version = "1.14.12"
 description = "Vector Quantization - Pytorch"
 authors = [
     { name = "Phil Wang", email = "[email protected]" }

vector_quantize_pytorch/lookup_free_quantization.py

@@ -48,14 +48,6 @@ def log(t, eps = 1e-5):
 def entropy(prob):
     return (-prob * log(prob)).sum(dim=-1)
 
-# distance
-
-def euclidean_distance_squared(x, y):
-    x2 = reduce(x ** 2, '... n d -> ... n 1', 'sum')
-    y2 = reduce(y ** 2, 'n d -> n', 'sum')
-    xy = einsum('... i d, j d -> ... i j', x, y) * -2
-    return x2 + xy + y2
-
 # class
 
 class LFQ(Module):
@@ -226,7 +218,8 @@ def forward(
         # entropy aux loss
 
         if self.training:
-            distance = euclidean_distance_squared(original_input, self.codebook)
+            # the same as euclidean distance up to a constant
+            distance = -2 * einsum('... i d, j d -> ... i j', original_input, self.codebook)
 
             prob = (-distance * inv_temperature).softmax(dim = -1)