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.)

Author: lucidrains

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,23 +12,29 @@ 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
 
 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