[MultiThreshold] clean up layout handling in MT custom op

should now support arbitrary layouts, but remaining backwards compatible
with the data_layout attribute may be a problem

Author: Yaman Umuroglu

src/qonnx/custom_op/general/multithreshold.py

@@ -33,16 +33,18 @@
 from qonnx.custom_op.base import CustomOp
 
 
-def multithreshold(v, thresholds, out_scale=None, out_bias=None):
+def multithreshold(v, thresholds, out_scale=None, out_bias=None, channels_last=False):
     """Given a set of threshold values t={t_0, t_1 ... t_n} the successive
     thresholding maps any real number x to an integer in the interval [0, n],
     where the returned integer is the number of thresholds x is greater than
     or equal to.
 
     The output tensor will be scaled by out_scale and biased by out_bias."""
+    # if channels_last=False:
     # the inputs are expected to be in the shape (N,C,_) where:
     #  C is the number of channels
     #  _ represents any (including zero) number of spatial dims
+    # if channels_last=True, expected input shape is (N,_,C)
     # the thresholds are expected to be in the shape (C, B) where
     # C : Number of channels (must be the same value as C in input tensor
     #     or 1 if all channels use the same threshold value)
@@ -52,16 +54,20 @@ def multithreshold(v, thresholds, out_scale=None, out_bias=None):
     # assert threshold shape - threshold channels must be either equal
     # to input channels, or be a single global scalar
     is_global_threshold = thresholds.shape[0] == 1
+    channel_axis = -1 if channels_last else 1
     assert (
-        v.shape[1] == thresholds.shape[0]
+        v.shape[channel_axis] == thresholds.shape[0]
     ) or is_global_threshold, """"Threshold
     shape incorrect"""
-    # starting assumption: input tensor is in NC_ layout
-    # (where _ can be any number of spatial dims)
-    # get the input tensor into right shape to use numpy broadcasting
-    # move the channels axis to the last position
-    # NC_ -> N_C
-    vm = np.moveaxis(v, source=1, destination=-1)
+    if not channels_last:
+        # starting assumption: input tensor is in NC_ layout
+        # (where _ can be any number of spatial dims)
+        # get the input tensor into right shape to use numpy broadcasting
+        # move the channels axis to the last position
+        # NC_ -> N_C
+        vm = np.moveaxis(v, source=1, destination=-1)
+    else:
+        vm = v
     # add a dummy dimension at the end of the input tensor
     # (for broadcasting against the thresholds)
     # N_C -> N_C1
@@ -72,8 +78,9 @@ def multithreshold(v, thresholds, out_scale=None, out_bias=None):
     # replace last axis by count of nonzero values (True)
     # N_CT -> N_C
     ret = np.count_nonzero(cmp, axis=-1)
-    # finally, move the channels axis back to index 1
-    ret = np.moveaxis(ret, source=-1, destination=1)
+    if not channels_last:
+        # move the channels axis back to index 1
+        ret = np.moveaxis(ret, source=-1, destination=1)
     assert ret.shape == v.shape, "Shape changed during thresholding!"
 
     if out_scale is None:
@@ -123,51 +130,11 @@ def execute_node(self, context, graph):
         out_bias = self.get_nodeattr("out_bias")
         # transpose input if NHWC data layout is chosen
         data_layout = self.get_nodeattr("data_layout")
-        if data_layout == "NHWC":
-            if v.ndim == 4:
-                # NHWC -> NCHW
-                v = np.transpose(v, (0, 3, 1, 2))
-            elif v.ndim == 2:
-                # no HW dimension means NHWC and NCHW layouts are equivalent
-                pass
-            else:
-                raise Exception("Unknown data_layout and input ndim" " combination for MultiThreshold.")
-
-        # Remember whether the shape has been modified to handle 1d or 3d data
-        # layouts
-        orig_shape = None
-        # If the input tensor has dimensions not covered by the NC or NCWH data
-        # layouts, the shape needs to be adapted such that it can be handled by
-        # multithreshold.
-        # TODO: Seems like a rather sketchy solution to support arbitrary data
-        #  layouts. This does not even validate the assumption of channel last
-        #  layout.
-        if v.ndim not in {2, 4}:
-            # Remember the original shape to be restored later
-            orig_shape = v.shape
-            # Assume last dimension to be the channel dimension C and reshape
-            # into NC layout which is supported by multithreshold
-            v = v.reshape((-1, v.shape[-1]))
-
+        channels_last = True if data_layout[-1] == "C" else False
         # calculate output
-        output = multithreshold(v, thresholds, out_scale, out_bias)
-        # setting context according to output
-        if data_layout == "NHWC":
-            if output.ndim == 4:
-                # NCHW -> NHWC
-                output = np.transpose(output, (0, 2, 3, 1))
-            elif output.ndim == 2:
-                # no HW dimension means NHWC and NCHW layouts are equivalent
-                pass
-            else:
-                raise Exception("Unknown data_layout and output ndim" " combination for MultiThreshold.")
-
-        # If the shape has been modified to support arbitrary layouts, restore
-        # the original shape
-        # TODO: Part of the rather sketchy solution above.
-        if orig_shape is not None:
-            output = output.reshape(orig_shape)
-
+        orig_shape = v.shape
+        output = multithreshold(v, thresholds, out_scale, out_bias, channels_last)
+        assert output.shape == orig_shape, "Shape changed during thresholding!"
         context[node.output[0]] = output
 
     def verify_node(self):