Complete xlayers.rearrange

xautodl/xlayers/super_rearrange.py

@@ -8,10 +8,14 @@ import torch.nn as nn
 import torch.nn.functional as F
 
 import math
+import numpy as np
+import itertools
+import functools
+from collections import OrderedDict
 from typing import Optional, Callable
 
 from xautodl import spaces
-from .misc_utils import ParsedExpression
+from .misc_utils import ParsedExpression, AnonymousAxis
 from .super_module import SuperModule
 from .super_module import IntSpaceType
 from .super_module import BoolSpaceType
@@ -31,11 +35,133 @@ class SuperReArrange(SuperModule):
         left, right = pattern.split("->")
         left = ParsedExpression(left)
         right = ParsedExpression(right)
+        difference = set.symmetric_difference(left.identifiers, right.identifiers)
+        if difference:
+            raise ValueError(
+                "Identifiers only on one side of expression (should be on both): {}".format(
+                    difference
+                )
+            )
 
-        import pdb
+        # parsing all dimensions to find out lengths
+        axis_name2known_length = OrderedDict()
+        for composite_axis in left.composition:
+            for axis_name in composite_axis:
+                if isinstance(axis_name, AnonymousAxis):
+                    axis_name2known_length[axis_name] = axis_name.value
+                else:
+                    axis_name2known_length[axis_name] = None
+        for axis_name in right.identifiers:
+            if axis_name not in axis_name2known_length:
+                if isinstance(axis_name, AnonymousAxis):
+                    axis_name2known_length[axis_name] = axis_name.value
+                else:
+                    axis_name2known_length[axis_name] = None
 
-        pdb.set_trace()
-        print("-")
+        axis_name2position = {
+            name: position for position, name in enumerate(axis_name2known_length)
+        }
+        for elementary_axis, axis_length in axes_lengths:
+            if not ParsedExpression.check_axis_name(elementary_axis):
+                raise ValueError("Invalid name for an axis", elementary_axis)
+            if elementary_axis not in axis_name2known_length:
+                raise ValueError(
+                    "Axis {} is not used in transform".format(elementary_axis)
+                )
+            axis_name2known_length[elementary_axis] = axis_length
+
+        input_composite_axes = []
+        # some of shapes will be inferred later - all information is prepared for faster inference
+        for composite_axis in left.composition:
+            known = {
+                axis
+                for axis in composite_axis
+                if axis_name2known_length[axis] is not None
+            }
+            unknown = {
+                axis for axis in composite_axis if axis_name2known_length[axis] is None
+            }
+            if len(unknown) > 1:
+                raise ValueError("Could not infer sizes for {}".format(unknown))
+            assert len(unknown) + len(known) == len(composite_axis)
+            input_composite_axes.append(
+                (
+                    [axis_name2position[axis] for axis in known],
+                    [axis_name2position[axis] for axis in unknown],
+                )
+            )
+
+        axis_position_after_reduction = {}
+        for axis_name in itertools.chain(*left.composition):
+            if axis_name in right.identifiers:
+                axis_position_after_reduction[axis_name] = len(
+                    axis_position_after_reduction
+                )
+
+        result_axes_grouping = []
+        for composite_axis in right.composition:
+            result_axes_grouping.append(
+                [axis_name2position[axis] for axis in composite_axis]
+            )
+
+        ordered_axis_right = list(itertools.chain(*right.composition))
+        axes_permutation = tuple(
+            axis_position_after_reduction[axis]
+            for axis in ordered_axis_right
+            if axis in left.identifiers
+        )
+        #
+        self.input_composite_axes = input_composite_axes
+        self.output_composite_axes = result_axes_grouping
+        self.elementary_axes_lengths = list(axis_name2known_length.values())
+        self.axes_permutation = axes_permutation
+
+    @functools.lru_cache(maxsize=1024)
+    def reconstruct_from_shape(self, shape):
+        if len(shape) != len(self.input_composite_axes):
+            raise ValueError(
+                "Expected {} dimensions, got {}".format(
+                    len(self.input_composite_axes), len(shape)
+                )
+            )
+        axes_lengths = list(self.elementary_axes_lengths)
+        for input_axis, (known_axes, unknown_axes) in enumerate(
+            self.input_composite_axes
+        ):
+            length = shape[input_axis]
+            known_product = 1
+            for axis in known_axes:
+                known_product *= axes_lengths[axis]
+            if len(unknown_axes) == 0:
+                if (
+                    isinstance(length, int)
+                    and isinstance(known_product, int)
+                    and length != known_product
+                ):
+                    raise ValueError(
+                        "Shape mismatch, {} != {}".format(length, known_product)
+                    )
+            else:
+                if (
+                    isinstance(length, int)
+                    and isinstance(known_product, int)
+                    and length % known_product != 0
+                ):
+                    raise ValueError(
+                        "Shape mismatch, can't divide axis of length {} in chunks of {}".format(
+                            length, known_product
+                        )
+                    )
+
+                (unknown_axis,) = unknown_axes
+                axes_lengths[unknown_axis] = length // known_product
+        # at this point all axes_lengths are computed (either have values or variables, but not Nones)
+        final_shape = []
+        for output_axis, grouping in enumerate(self.output_composite_axes):
+            lengths = [axes_lengths[elementary_axis] for elementary_axis in grouping]
+            final_shape.append(int(np.prod(lengths)))
+        axes_reordering = self.axes_permutation
+        return axes_lengths, axes_reordering, final_shape
 
     @property
     def abstract_search_space(self):
@@ -46,10 +172,13 @@ class SuperReArrange(SuperModule):
         self.forward_raw(input)
 
     def forward_raw(self, input: torch.Tensor) -> torch.Tensor:
-        import pdb
-
-        pdb.set_trace()
-        raise NotImplementedError
+        init_shape, axes_reordering, final_shape = self.reconstruct_from_shape(
+            tuple(input.shape)
+        )
+        tensor = torch.reshape(input, init_shape)
+        tensor = tensor.permute(axes_reordering)
+        tensor = torch.reshape(tensor, final_shape)
+        return tensor
 
     def extra_repr(self) -> str:
         params = repr(self._pattern)