+ k = 0
+ for d in range(4):
+ while True:
+ ii, jj = (
+ torch.randint(self.height, (1,)).item(),
+ torch.randint(self.width, (1,)).item(),
+ )
+ e = 0
+ for i in range(S):
+ for j in range(S):
+ if (
+ ii + i >= self.height
+ or jj + j >= self.width
+ or (
+ f_X[i + i0, j + j0] == c[d]
+ and X[ii + i, jj + j] > 0
+ )
+ ):
+ e = 1
+ if e == 0:
+ break
+ for i in range(S):
+ for j in range(S):
+ if f_X[i + i0, j + j0] == c[d]:
+ X[ii + i, jj + j] = c[d]
+
+ def TOO_MESSY_task_islands(self, A, f_A, B, f_B):
+ c = torch.randperm(self.nb_colors - 1)[:2] + 1
+ for X, f_X in [(A, f_A), (B, f_B)]:
+ if not hasattr(self, "cache_islands") or len(self.cache_islands) == 0:
+ self.cache_islands = list(
+ grow_islands(
+ 1000,
+ self.height,
+ self.width,
+ nb_seeds=self.height * self.width // 20,
+ nb_iterations=self.height * self.width // 2,
+ )
+ )
+
+ A = self.cache_islands.pop()
+
+ while True:
+ i, j = (
+ torch.randint(self.height // 2, (1,)).item(),
+ torch.randint(self.width // 2, (1,)).item(),
+ )
+ if A[i, j] > 0:
+ break
+
+ X[...] = (A > 0) * c[0]
+ f_X[...] = (A == A[i, j]) * c[1] + ((A > 0) & (A != A[i, j])) * c[0]
+ f_X[i, j] = X[i, j]
+ X[i, j] = c[1]
+
+ # @torch.compile
+ def TOO_HARD_task_stack(self, A, f_A, B, f_B):
+ N = 5
+ c = torch.randperm(self.nb_colors - 1)[:N] + 1
+ for X, f_X in [(A, f_A), (B, f_B)]:
+ i1, j1, i2, j2 = (
+ self.height // 2 - 1,
+ self.width // 2 - 1,
+ self.height // 2 + 1,
+ self.width // 2 + 1,
+ )
+ op = torch.tensor((0, 1, 2, 3) * 4)
+ op = op[torch.randperm(op.size(0))[:9]]
+ for q in range(op.size(0)):
+ u = 3 * (q // 3)
+ v = 3 * (q % 3)
+ d = c[torch.randint(N, (1,)).item()]
+ # X[u+1,v+1]=d
+ if op[q] == 0: # right
+ X[u : u + 3, v + 2] = d
+ elif op[q] == 1: # let
+ X[u : u + 3, v] = d
+ elif op[q] == 2: # bottom
+ X[u + 2, v : v + 3] = d
+ elif op[q] == 3: # top
+ X[u, v : v + 3] = d
+
+ if q == 0:
+ f_X[i1:i2, j1:j2] = d
+ elif op[q] == 0: # right
+ f_X[i1:i2, j2] = d
+ j2 += 1
+ elif op[q] == 1: # let
+ j1 -= 1
+ f_X[i1:i2, j1] = d
+ elif op[q] == 2: # bottom
+ f_X[i2, j1:j2] = d
+ i2 += 1
+ elif op[q] == 3: # top
+ i1 -= 1
+ f_X[i1, j1:j2] = d
+
+ def randint(self, *m):
+ m = torch.tensor(m)
+ return (torch.rand(m.size()) * m).long()
+
+ def TOO_HARD_task_matrices(self, A, f_A, B, f_B):
+ N = 6
+ c = torch.randperm(self.nb_colors - 1)[:N] + 1
+
+ for X, f_X in [(A, f_A), (B, f_B)]:
+ M1 = torch.randint(2, (5, 5))
+ M2 = torch.randint(2, (5, 5))
+ P = M1 @ M2
+ for i in range(5):
+ for j in range(5):
+ X[i, j] = c[M1[i, j]]
+ X[i, j + 5] = c[M2[i, j]]
+ f_X[i, j] = c[M1[i, j]]
+ f_X[i, j + 5] = c[M2[i, j]]
+ f_X[i + 5, j + 5] = c[P[i, j]]
+
+ def TOO_HARD_task_compute(self, A, f_A, B, f_B):
+ N = 6
+ c = torch.randperm(self.nb_colors - 1)[:N] + 1
+ for X, f_X in [(A, f_A), (B, f_B)]:
+ v = torch.randint((self.width - 1) // 2, (N,)) + 1
+ chain = torch.randperm(N)
+ eq = []
+ for i in range(chain.size(0) - 1):
+ i1, i2 = chain[i], chain[i + 1]
+ v1, v2 = v[i1], v[i2]
+ k = torch.arange(self.width // 2) + 1
+ d = ((k[None, :] * v1 - k[:, None] * v2) == 0).nonzero() + 1
+ d = d[torch.randint(d.size(0), (1,)).item()]
+ w1, w2 = d
+ eq.append((c[i1], w1, c[i2], w2))
+
+ ii = torch.randperm(self.height - 2)[: len(eq)]
+
+ for k, x in enumerate(eq):
+ i = ii[k]
+ c1, w1, c2, w2 = x
+ s = torch.randint(self.width - (w1 + w2) + 1, (1,)).item()
+ X[i, s : s + w1] = c1
+ X[i, s + w1 : s + w1 + w2] = c2
+ f_X[i, s : s + w1] = c1
+ f_X[i, s + w1 : s + w1 + w2] = c2
+
+ i1, i2 = torch.randperm(N)[:2]
+ v1, v2 = v[i1], v[i2]
+ k = torch.arange(self.width // 2) + 1
+ d = ((k[None, :] * v1 - k[:, None] * v2) == 0).nonzero() + 1
+ d = d[torch.randint(d.size(0), (1,)).item()]
+ w1, w2 = d
+ c1, c2 = c[i1], c[i2]
+ s = 0 # torch.randint(self.width - (w1 + w2) + 1, (1,)).item()
+ i = self.height - 1
+ X[i, s : s + w1] = c1
+ X[i, s + w1 : s + w1 + 1] = c2
+ f_X[i, s : s + w1] = c1
+ f_X[i, s + w1 : s + w1 + w2] = c2
+
+ # @torch.compile
+ # [ai1,ai2] [bi1,bi2]
+ def task_contact(self, A, f_A, B, f_B):
+ def rec_dist(a, b):
+ ai1, aj1, ai2, aj2 = a
+ bi1, bj1, bi2, bj2 = b
+ v = max(ai1 - bi2, bi1 - ai2)
+ h = max(aj1 - bj2, bj1 - aj2)
+ return min(max(v, 0) + max(h + 1, 0), max(v + 1, 0) + max(h, 0))
+
+ nb_rec = 3
+ c = torch.randperm(self.nb_colors - 1)[:nb_rec] + 1
+ for X, f_X in [(A, f_A), (B, f_B)]:
+ while True:
+ r = self.rec_coo(nb_rec, prevent_overlap=True)
+ d = [rec_dist(r[0], r[k]) for k in range(nb_rec)]
+ if min(d[1:]) == 0:
+ break
+
+ for n in range(nb_rec):
+ i1, j1, i2, j2 = r[n]
+ X[i1:i2, j1:j2] = c[n]
+ f_X[i1:i2, j1:j2] = c[n]
+ if d[n] == 0:
+ f_X[i1, j1:j2] = c[0]
+ f_X[i2 - 1, j1:j2] = c[0]
+ f_X[i1:i2, j1] = c[0]
+ f_X[i1:i2, j2 - 1] = c[0]
+
+ # @torch.compile
+ # [ai1,ai2] [bi1,bi2]
+ def task_corners(self, A, f_A, B, f_B):
+ polarity = torch.randint(2, (1,)).item()
+ nb_rec = 3
+ c = torch.randperm(self.nb_colors - 1)[:nb_rec] + 1
+ for X, f_X in [(A, f_A), (B, f_B)]:
+ r = self.rec_coo(nb_rec, prevent_overlap=True)
+
+ for n in range(nb_rec):
+ i1, j1, i2, j2 = r[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)
+ d = dd[1:-1, 1:-1]
+ m = (X > 0).long()
+ d[i, j] = 0
+ e = d.clone()
+ while True:
+ e[...] = d
+ d[...] = (
+ d.min(dd[:-2, 1:-1] + 1)
+ .min(dd[2:, 1:-1] + 1)
+ .min(dd[1:-1, :-2] + 1)
+ .min(dd[1:-1, 2:] + 1)
+ )
+ d[...] = (1 - m) * d + m * self.height * self.width
+ if e.equal(d):
+ break
+
+ return d
+
+ # @torch.compile
+ def task_path(self, A, f_A, B, f_B):
+ nb_rec = 2
+ c = torch.randperm(self.nb_colors - 1)[: nb_rec + 2] + 1
+ for X, f_X in [(A, f_A), (B, f_B)]:
+ while True:
+ X[...] = 0
+ f_X[...] = 0
+
+ r = self.rec_coo(nb_rec, prevent_overlap=True)
+ for n in range(nb_rec):
+ i1, j1, i2, j2 = r[n]
+ X[i1:i2, j1:j2] = c[n]
+ f_X[i1:i2, j1:j2] = c[n]
+
+ i1, i2 = torch.randint(self.height, (2,))
+ j1, j2 = torch.randint(self.width, (2,))
+ if (
+ abs(i1 - i2) + abs(j1 - j2) > 2
+ and X[i1, j1] == 0
+ and X[i2, j2] == 0
+ ):
+ d2 = self.compdist(X, i2, j2)
+ d = self.compdist(X, i1, j1)
+
+ if d2[i1, j1] < 2 * self.width:
+ break
+
+ m = ((d + d2) == d[i2, j2]).long()
+ f_X[...] = m * c[-1] + (1 - m) * f_X
+
+ X[i1, j1] = c[-2]
+ X[i2, j2] = c[-2]
+ f_X[i1, j1] = c[-2]
+ f_X[i2, j2] = c[-2]
+
+ # @torch.compile
+ def task_fill(self, A, f_A, B, f_B):
+ nb_rec = 3
+ c = torch.randperm(self.nb_colors - 1)[: nb_rec + 1] + 1
+ for X, f_X in [(A, f_A), (B, f_B)]:
+ accept_full = torch.rand(1) < 0.5
+
+ while True:
+ X[...] = 0
+ f_X[...] = 0
+
+ r = self.rec_coo(nb_rec, prevent_overlap=True)
+ for n in range(nb_rec):
+ i1, j1, i2, j2 = r[n]
+ X[i1:i2, j1:j2] = c[n]
+ f_X[i1:i2, j1:j2] = c[n]
+
+ while True:
+ i, j = (
+ torch.randint(self.height, (1,)).item(),
+ torch.randint(self.width, (1,)).item(),
+ )
+ if X[i, j] == 0:
+ break