From abdc7be721e3c0dcd8b4f08bd6b87f67f73f8188 Mon Sep 17 00:00:00 2001
From: lucidrains <lucidrains@gmail.com>
Date: Thu, 24 Apr 2025 08:25:40 -0700
Subject: [PATCH 1/3] begin to test out a vector quantizer with elements of
 genetic algorithms. share the idea

---
 vector_quantize_pytorch/evo_vq.py | 0
 1 file changed, 0 insertions(+), 0 deletions(-)
 create mode 100644 vector_quantize_pytorch/evo_vq.py

diff --git a/vector_quantize_pytorch/evo_vq.py b/vector_quantize_pytorch/evo_vq.py
new file mode 100644
index 0000000..e69de29

From d93053897eef0085ee0231717cc41c739283680b Mon Sep 17 00:00:00 2001
From: Phil Wang <lucidrains@gmail.com>
Date: Sat, 26 Apr 2025 06:02:14 -0700
Subject: [PATCH 2/3] benefits need to be seen with a pop of size 3

---
 vector_quantize_pytorch/evo_vq.py | 41 +++++++++++++++++++++++++++++++
 1 file changed, 41 insertions(+)

diff --git a/vector_quantize_pytorch/evo_vq.py b/vector_quantize_pytorch/evo_vq.py
index e69de29..6207beb 100644
--- a/vector_quantize_pytorch/evo_vq.py
+++ b/vector_quantize_pytorch/evo_vq.py
@@ -0,0 +1,41 @@
+import torch
+from torch import cat
+
+# helpers
+
+def exists(v):
+    return v is not None
+
+def default(v, d):
+    return v if exists(v) else d
+
+# evolution - start with the most minimal, a population of 3
+# 1 is natural selected out, the other 2 performs crossover
+
+def select_and_crossover(
+    codes,   # Float[3 ...]
+    fitness, # Float[3]
+):
+    assert codes.shape[0] == fitness.shape[0] == 3
+
+    # selection
+
+    top2 = fitness.topk(2, dim = -1).indices
+    codes = codes[top2]
+
+    # crossover
+
+    child = codes.mean(dim = 0, keepdim = True)
+    codes = cat((codes, child))
+
+    return codes
+
+# class
+
+class EvoVQ(Module):
+    def __init__(self):
+        super().__init__()
+        raise NotImplementedError
+
+    def forward(self, x):
+        raise NotImplementedError

From e957775af70f3f5456cf5ab794cf1631e0957569 Mon Sep 17 00:00:00 2001
From: Phil Wang <lucidrains@gmail.com>
Date: Sat, 26 Apr 2025 06:23:20 -0700
Subject: [PATCH 3/3] expand to any multiple of 3 for pop size, select out
 bottom third and replenish with rest of the paired couples with approximate
 fitness uniform strategy (Hutter et al.)

---
 vector_quantize_pytorch/evo_vq.py | 17 +++++++++++++----
 1 file changed, 13 insertions(+), 4 deletions(-)

diff --git a/vector_quantize_pytorch/evo_vq.py b/vector_quantize_pytorch/evo_vq.py
index 6207beb..4a82067 100644
--- a/vector_quantize_pytorch/evo_vq.py
+++ b/vector_quantize_pytorch/evo_vq.py
@@ -1,5 +1,8 @@
 import torch
 from torch import cat
+from torch.nn import Module
+
+from einops import reduce
 
 # helpers
 
@@ -9,6 +12,9 @@ def exists(v):
 def default(v, d):
     return v if exists(v) else d
 
+def divisible_by(num, den):
+    return (num % den) == 0
+
 # evolution - start with the most minimal, a population of 3
 # 1 is natural selected out, the other 2 performs crossover
 
@@ -16,16 +22,19 @@ def select_and_crossover(
     codes,   # Float[3 ...]
     fitness, # Float[3]
 ):
-    assert codes.shape[0] == fitness.shape[0] == 3
+    pop_size = codes.shape[0]
+    assert pop_size == fitness.shape[0]
+    assert divisible_by(pop_size, 3)
 
     # selection
 
-    top2 = fitness.topk(2, dim = -1).indices
-    codes = codes[top2]
+    sorted_indices = fitness.sort().indices
+    selected = sorted_indices[(pop_size // 3):] # bottom third wins darwin awards
+    codes = codes[selected]
 
     # crossover
 
-    child = codes.mean(dim = 0, keepdim = True)
+    child = reduce(codes, '(two paired) ... -> paired ...', 'mean', two = 2)
     codes = cat((codes, child))
 
     return codes