+ def task_count(self, A, f_A, B, f_B):
+ N = torch.randint(3, (1,)) + 1
+ c = torch.randperm(len(self.colors) - 1)[:N] + 1
+ for X, f_X in [(A, f_A), (B, f_B)]:
+ nb = torch.randint(self.width, (3,)) + 1
+ k = torch.randperm(self.height * self.width)[: nb.sum()]
+ p = 0
+ for n in range(N):
+ for m in range(nb[n]):
+ i, j = k[p] % self.height, k[p] // self.height
+ X[i, j] = c[n]
+ f_X[n, m] = c[n]
+ p += 1
+