Update.

diff --git a/reasoning.py b/reasoning.py
index c442947..57e8056 100755 (executable)
--- a/reasoning.py
+++ b/reasoning.py
@@ -171,11 +171,11 @@ class Reasoning(problem.Problem):
 
     # That's quite a tensorial spaghetti mess to sample
     # non-overlapping rectangles quickly, but made the generation of
-    # 100k samples from 1h50 with a lame pure python code to 4min with
-    # this one.
+    # 100k samples go from 1h50 with a lame pure python code to 3min30s
+    # with this one.
     def rec_coo(self, x, n, min_height=3, min_width=3):
         K = 3
-        N = 1000
+        N = 200
 
         while True:
             v = (
@@ -322,19 +322,42 @@ class Reasoning(problem.Problem):
         di, dj = torch.randint(2, (2,)) * 2 - 1
         N = 3
         c = torch.randperm(len(self.colors) - 1)[: 2 * N] + 1
-        direction = torch.randint(2, (1,))
+        direction = torch.randint(4, (1,))
         for X, f_X in [(A, f_A), (B, f_B)]:
             r = self.rec_coo(X, N)
             for n in range(N):
                 i1, j1, i2, j2 = r[n]
-                i = (i1 + i2) // 2
                 X[i1:i2, j1:j2] = c[2 * n]
-                X[i : i + 1, j1:j2] = c[2 * n + 1]
                 f_X[i1:i2, j1:j2] = c[2 * n]
-                if n == N - 1:
-                    f_X[i:i2, j1:j2] = c[2 * n + 1]
-                else:
-                    f_X[i : i + 1, j1:j2] = c[2 * n + 1]
+                # Not my proudest moment
+                if direction == 0:
+                    i = (i1 + i2) // 2
+                    X[i : i + 1, j1:j2] = c[2 * n + 1]
+                    if n == N - 1:
+                        f_X[i:i2, j1:j2] = c[2 * n + 1]
+                    else:
+                        f_X[i : i + 1, j1:j2] = c[2 * n + 1]
+                elif direction == 1:
+                    i = (i1 + i2 - 1) // 2
+                    X[i : i + 1, j1:j2] = c[2 * n + 1]
+                    if n == N - 1:
+                        f_X[i1 : i + 1, j1:j2] = c[2 * n + 1]
+                    else:
+                        f_X[i : i + 1, j1:j2] = c[2 * n + 1]
+                elif direction == 2:
+                    j = (j1 + j2) // 2
+                    X[i1:i2, j : j + 1] = c[2 * n + 1]
+                    if n == N - 1:
+                        f_X[i1:i2, j:j2] = c[2 * n + 1]
+                    else:
+                        f_X[i1:i2, j : j + 1] = c[2 * n + 1]
+                elif direction == 3:
+                    j = (j1 + j2 - 1) // 2
+                    X[i1:i2, j : j + 1] = c[2 * n + 1]
+                    if n == N - 1:
+                        f_X[i1:i2, j1 : j + 1] = c[2 * n + 1]
+                    else:
+                        f_X[i1:i2, j : j + 1] = c[2 * n + 1]
 
     def task_frame(self, A, f_A, B, f_B):
         N = 3