Update

diff --git a/beaver.py b/beaver.py
index 8fe9a9b..69116ea 100755 (executable)
--- a/beaver.py
+++ b/beaver.py
@@ -135,24 +135,33 @@ def generation_order(x, fixed_len):
     if args.random_regression_order:
         order = torch.rand(x.size(), device=x.device)
         order[:, :fixed_len] = torch.linspace(-2, -1, fixed_len, device=order.device)
-        return order.sort(1).indices
+        order = order.sort(1).indices
     else:
-        return (
+        order = (
             torch.arange(x.size(1), device=x.device).unsqueeze(0).expand(x.size(0), -1)
         )
+    return order
 
 
-def shuffle(x, order, reorder=False):
-    if x.dim() == 3:
-        order = order.unsqueeze(-1).expand(-1, -1, x.size(-1))
-    if reorder:
-        y = x.new(x.size())
-        y.scatter_(1, order, x)
-        return y
+def reorder(x, order, back=False):  # x is NxTxD1x...xDk, order is NxT'
+    u = x.reshape(x.size()[:2] + (-1,))
+    order = order.unsqueeze(-1).expand(-1, -1, u.size(-1))
+    if back:
+        v = u.new(u.size())
+        v.scatter_(1, order, u)
     else:
-        return x.gather(1, order)
+        v = u.gather(1, order)
+    v = v.reshape(v.size()[:2] + x.size()[2:])
+    return v
 
 
+def shuffle(x, fixed_len):
+    order = generation_order(x, fixed_len)
+    return reorder(x, order), order
+
+
+######################################################################
+
 # ar_mask is a Boolean matrix of same shape as input, with 1s on the
 # tokens that should be generated
 
@@ -186,8 +195,7 @@ def compute_perplexity(model, split="train"):
 
         for input in task.batches(split=split):
             input = input.to(device)
-            order = generation_order(input, task.height * task.width)
-            input = shuffle(input, order)
+            input, order = shuffle(input, task.height * task.width)
             output = model(mygpt.BracketedSequence(input), order=order).x
             loss = F.cross_entropy(output.transpose(1, 2), input)
             acc_loss += loss.item() * input.size(0)
@@ -252,10 +260,9 @@ def oneshot(gpt, task):
 
         acc_train_loss, nb_train_samples = 0, 0
         for mazes, policies in task.policy_batches(split="train"):
-            order = generation_order(mazes, task.height * task.width)
-            x = shuffle(mazes, order)
+            x, order = shuffle(mazes, task.height * task.width)
             x = gpt(mygpt.BracketedSequence(x), mode=args.oneshot_input, order=order).x
-            output_gpt = shuffle(x, order, reorder=True)
+            output_gpt = reorder(x, order, back=True)
             output = model(output_gpt)
 
             loss = compute_loss(mazes, output, policies, task.height, task.width)
@@ -268,10 +275,9 @@ def oneshot(gpt, task):
 
         acc_test_loss, nb_test_samples = 0, 0
         for mazes, policies in task.policy_batches(split="test"):
-            order = generation_order(mazes, task.height * task.width)
-            x = shuffle(mazes, order)
+            x, order = shuffle(mazes, task.height * task.width)
             x = gpt(mygpt.BracketedSequence(x), mode=args.oneshot_input, order=order).x
-            output_gpt = shuffle(x, order, reorder=True)
+            output_gpt = reorder(x, order, back=True)
             output = model(output_gpt)
             loss = compute_loss(mazes, output, policies, task.height, task.width)
             acc_test_loss += loss.item() * mazes.size(0)
@@ -284,10 +290,9 @@ def oneshot(gpt, task):
         # -------------------
         mazes = task.test_input[:32, : task.height * task.width]
         policies = task.test_policies[:32]
-        order = generation_order(mazes, task.height * task.width)
-        x = shuffle(mazes, order)
+        x, order = shuffle(mazes, task.height * task.width)
         x = gpt(mygpt.BracketedSequence(x), mode=args.oneshot_input, order=order).x
-        output_gpt = shuffle(x, order, reorder=True)
+        output_gpt = reorder(x, order, back=True)
         output = model(output_gpt)
         if args.oneshot_output == "policy":
             targets = policies.permute(0, 2, 1)
@@ -426,11 +431,11 @@ class TaskMaze(Task):
             ar_mask = result.new_zeros(result.size())
             ar_mask[:, self.height * self.width :] = 1
             result *= 1 - ar_mask
-            order = generation_order(result, self.height * self.width)
+            result, order = shuffle(result, self.height * self.width)
             masked_inplace_autoregression(
                 model, self.batch_size, result, ar_mask, order=order
             )
-            result = shuffle(result, order, reorder=True)
+            result = reorder(result, order, back=True)
             mazes, paths = self.seq2map(result)
             nb_correct += maze.path_correctness(mazes, paths).long().sum()
             nb_total += mazes.size(0)
@@ -608,8 +613,7 @@ for n_epoch in range(nb_epochs_finished, args.nb_epochs):
 
     for input in task.batches(split="train"):
         input = input.to(device)
-        order = generation_order(input, task.height * task.width)
-        input = shuffle(input, order)
+        input, order = shuffle(input, task.height * task.width)
         output = model(mygpt.BracketedSequence(input), order=order).x
         loss = F.cross_entropy(output.transpose(1, 2), input)
         acc_train_loss += loss.item() * input.size(0)