+ def task_color_grow(self, A, f_A, B, f_B):
+ 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,))
+ 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]
+ X[i1 : (i1 + i2) // 2, j1:j2] = c[2 * n]
+ f_X[i1 : (i1 + i2) // 2, j1:j2] = c[2 * n]
+ X[(i1 + i2) // 2 : (i1 + i2) // 2 + 1, j1:j2] = c[2 * n + 1]
+ if n == N - 1:
+ f_X[(i1 + i2) // 2 : i2, j1:j2] = c[2 * n + 1]