Update.

diff --git a/grids.py b/grids.py
index f747c47..3343066 100755 (executable)
--- a/grids.py
+++ b/grids.py
@@ -172,13 +172,11 @@ class Grids(problem.Problem):
         S = self.height * self.width
 
         mask = torch.tensor(mask, device=quizzes.device)
-        mask = mask[None, :, None].expand(1, 4, S + 1)
+        mask = mask[None, :, None].expand(1, 4, S + 1).clone()
         mask[:, :, 0] = 0
         mask = mask.reshape(1, -1).expand_as(quizzes)
-
         mask = mask * (torch.rand(mask.size(), device=mask.device) <= noise).long()
         random = torch.randint(self.nb_colors, mask.size())
-
         quizzes = mask * random + (1 - mask) * quizzes
 
         return quizzes
@@ -712,7 +710,9 @@ class Grids(problem.Problem):
                 X[i1:i2, j1:j2] = c[n]
                 f_X[i1:i2, j1:j2] = c[n]
                 if n < nb_rec - 1:
-                    f_X[i1, j1] = c[-1]
+                    for k in range(2):
+                        f_X[i1 + k, j1] = c[-1]
+                        f_X[i1, j1 + k] = c[-1]
 
     # @torch.compile
     def contact(self, X, i, j, q):
@@ -925,9 +925,11 @@ class Grids(problem.Problem):
                 X[i + i1 : i + i2, j + j1 : j + j2] = c[0]
                 f_X[2 * i1 : 2 * i2, 2 * j1 : 2 * j2] = c[0]
 
-            X[i, j] = c[1]
-            # f_X[0:2, 0:2] = c[1]
-            f_X[i, j] = c[1]
+            for k in range(2):
+                X[i + k, j] = c[1]
+                X[i, j + k] = c[1]
+                f_X[i + k, j] = c[1]
+                f_X[i, j + k] = c[1]
 
     # @torch.compile
     def task_symbols(self, A, f_A, B, f_B):
@@ -953,18 +955,28 @@ class Grids(problem.Problem):
             for Z in [X, f_X]:
                 for k in range(1, nb_rec):
                     Z[i[k] : i[k] + delta, j[k] : j[k] + delta] = c[k]
-                Z[i[0] + delta // 2 - 1, j[0] + delta // 2 - 1] = c[0]
-                Z[i[0] + delta // 2 - 1, j[0] + delta // 2 + 1] = c[0]
-                Z[i[0] + delta // 2 + 1, j[0] + delta // 2 - 1] = c[0]
-                Z[i[0] + delta // 2 + 1, j[0] + delta // 2 + 1] = c[0]
+                # Z[i[0] + delta // 2 - 1, j[0] + delta // 2 - 1] = c[0]
+                # Z[i[0] + delta // 2 - 1, j[0] + delta // 2 + 1] = c[0]
+                # Z[i[0] + delta // 2 + 1, j[0] + delta // 2 - 1] = c[0]
+                # Z[i[0] + delta // 2 + 1, j[0] + delta // 2 + 1] = c[0]
 
             # f_X[i[0] : i[0] + delta, j[0] : j[0] + delta] = c[q]
-            f_X[i[0] + delta // 2, j[0] + delta // 2] = c[q]
 
-            X[
+            # f_X[i[0] + delta // 2, j[0] + delta // 2] = c[q]
+            f_X[i[0] : i[0] + delta, j[0] : j[0] + delta] = c[q]
+
+            ii, jj = (
                 i[0] + delta // 2 + (i[q] - ai).sign().long(),
                 j[0] + delta // 2 + (j[q] - aj).sign().long(),
-            ] = c[nb_rec]
+            )
+
+            X[ii, jj] = c[nb_rec]
+            X[i[0] + delta // 2, jj] = c[nb_rec]
+            X[ii, j[0] + delta // 2] = c[nb_rec]
+
+            f_X[ii, jj] = c[nb_rec]
+            f_X[i[0] + delta // 2, jj] = c[nb_rec]
+            f_X[ii, j[0] + delta // 2] = c[nb_rec]
 
     # @torch.compile
     def task_isometry(self, A, f_A, B, f_B):
@@ -1353,13 +1365,18 @@ class Grids(problem.Problem):
 
             for n in range(nb_rec):
                 i1, j1, i2, j2 = r[n]
-                if polarity == 0:
-                    X[i1, j1] = c[n]
-                    X[i2 - 1, j2 - 1] = c[n]
-                else:
-                    X[i1, j2 - 1] = c[n]
-                    X[i2 - 1, j1] = c[n]
-                f_X[i1:i2, j1:j2] = c[n]
+                for k in range(2):
+                    if polarity == 0:
+                        X[i1 + k, j1] = c[n]
+                        X[i2 - 1 - k, j2 - 1] = c[n]
+                        X[i1, j1 + k] = c[n]
+                        X[i2 - 1, j2 - 1 - k] = c[n]
+                    else:
+                        X[i1 + k, j2 - 1] = c[n]
+                        X[i2 - 1 - k, j1] = c[n]
+                        X[i1, j2 - 1 - k] = c[n]
+                        X[i2 - 1, j1 + k] = c[n]
+                    f_X[i1:i2, j1:j2] = c[n]
 
     def compdist(self, X, i, j):
         dd = X.new_full((self.height + 2, self.width + 2), self.height * self.width)
@@ -1695,26 +1712,21 @@ if __name__ == "__main__":
     nb = 1000
 
     for t in [
-        grids.task_addition,
-        grids.task_bounce,
-        grids.task_compute,
-        grids.task_contact,
-        grids.task_corners,
-        grids.task_detect,
-        grids.task_fill,
-        grids.task_frame,
-        grids.task_grow,
-        grids.task_half_fill,
-        grids.task_islands,
-        grids.task_isometry,
-        grids.task_path,
-        grids.task_puzzle,
-        grids.task_replace_color,
-        grids.task_scale,
-        grids.task_stack,
+        # grids.task_bounce,
+        # grids.task_contact,
+        # grids.task_corners,
+        # grids.task_detect,
+        # grids.task_fill,
+        # grids.task_frame,
+        # grids.task_grow,
+        # grids.task_half_fill,
+        # grids.task_isometry,
+        # grids.task_path,
+        # grids.task_replace_color,
+        # grids.task_scale,
         grids.task_symbols,
-        grids.task_trajectory,
-        grids.task_translate,
+        # grids.task_trajectory,
+        # grids.task_translate,
     ]:
         # for t in [grids.task_path]:
         start_time = time.perf_counter()

diff --git a/quiz_machine.py b/quiz_machine.py
index a4ca60b..bfa7f97 100755 (executable)
--- a/quiz_machine.py
+++ b/quiz_machine.py
@@ -183,13 +183,14 @@ class QuizMachine:
 
         if self.prompt_noise > 0.0:
             for struct, mask in self.understood_structures:
-                i = self.problem.indices_select(quizzes=input, struct=struct)
-                input[i] = self.problem.inject_noise(
-                    input[i],
-                    self.prompt_noise,
-                    struct=struct,
-                    mask=(1 - k for k in mask),
-                )
+                i = self.problem.indices_select(quizzes=quizzes, struct=struct)
+                if i.any():
+                    quizzes[i] = self.problem.inject_noise(
+                        quizzes[i],
+                        self.prompt_noise,
+                        struct=struct,
+                        mask=tuple(1 - k for k in mask),
+                    )
 
         return quizzes, from_w