[MO] Add explicit broadcasting mode (#2077)

* [MO] Add explicit broadcasting mode

* corrected value broadcast, added nontrivial test cases

* added negative unit-tests

* added axes normalization and permute inputs

* corrected axes normalization

* Update model-optimizer/mo/utils/broadcasting.py

Co-authored-by: Evgenya Stepyreva <evgenya.stepyreva@intel.com>

model-optimizer/mo/ops/broadcast.py

@@ -18,7 +18,7 @@ from mo.graph.graph import Node, Graph
 from mo.graph.perm_inputs import PermuteInputs
 from mo.ops.op import Op
 from mo.utils.broadcasting import bi_directional_shape_broadcasting, uni_directional_shape_broadcasting, \
-    uni_directional_broadcasting, bi_directional_broadcasting
+    uni_directional_broadcasting, bi_directional_broadcasting, explicit_broadcasting, explicit_shape_broadcasting
 from mo.utils.error import Error
 
 
@@ -28,7 +28,7 @@ class Broadcast(Op):
         Inputs:
             [0] - tensor to be broadcasted
             [1] - shape to be broadcast to
-            [2] - optional axis parameter that which axis are allowed to be broadcasted
+            [2] - optional axes_mapping tensor
     """
 
     op = 'Broadcast'
@@ -69,6 +69,13 @@ class Broadcast(Op):
                 node.out_port(0).data.set_value(uni_directional_broadcasting(input_value, target_shape))
             elif node.mode == 'bidirectional':
                 node.out_port(0).data.set_value(bi_directional_broadcasting(input_value, target_shape))
+            elif node.mode == 'explicit':
+                axes_mapping = node.in_port(2).data.get_value()
+                assert axes_mapping  is not None, 'Broadcast(mode="explicit") with dynamic axes_mapping input ' \
+                                                  'is not supported. Node: `{}`'.format(node_name)
+                PermuteInputs().set_input_permutation(node.in_node(2), node, 'output:0', 'axis')
+                axes_mapping = node.in_port(2).data.get_value()
+                node.out_port(0).data.set_value(explicit_broadcasting(input_value, target_shape, axes_mapping))
             else:
                 raise Error('The node "{}" has unsupported mode "{}"'.format(node_name, node.mode))
         else:
@@ -76,5 +83,12 @@ class Broadcast(Op):
                 node.out_port(0).data.set_shape(uni_directional_shape_broadcasting(input_shape, target_shape))
             elif node.mode == 'bidirectional':
                 node.out_port(0).data.set_shape(bi_directional_shape_broadcasting(input_shape, target_shape))
+            elif node.mode == 'explicit':
+                axes_mapping = node.in_port(2).data.get_value()
+                assert axes_mapping is not None, 'Broadcast(mode="explicit") with dynamic axes_mapping input ' \
+                                                 'is not supported. Node: `{}`'.format(node_name)
+                axes_mapping = node.in_port(2).data.get_value()
+                new_shape,_ = explicit_shape_broadcasting(input_shape, target_shape, axes_mapping)
+                node.out_port(0).data.set_shape(new_shape)
             else:
                 raise Error('The node "{}" has unsupported mode "{}"'.format(node_name, node.mode))

model-optimizer/mo/ops/broadcast_test.py

@@ -0,0 +1,91 @@
+"""
+ Copyright (C) 2018-2020 Intel Corporation
+
+ Licensed under the Apache License, Version 2.0 (the "License");
+ you may not use this file except in compliance with the License.
+ You may obtain a copy of the License at
+
+      http://www.apache.org/licenses/LICENSE-2.0
+
+ Unless required by applicable law or agreed to in writing, software
+ distributed under the License is distributed on an "AS IS" BASIS,
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ See the License for the specific language governing permissions and
+ limitations under the License.
+"""
+
+import unittest
+
+import numpy as np
+from generator import generator, generate
+
+from mo.front.common.partial_infer.utils import int64_array
+from mo.graph.graph import Node
+from mo.ops.broadcast import Broadcast
+from mo.utils.unittest.graph import build_graph, valued_const_with_data, regular_op_with_empty_data, \
+    shaped_data
+
+
+@generator
+class BroadcastTest(unittest.TestCase):
+    @generate(*[
+        ([1], [3, 3], None, 'numpy', [[1, 1, 1], [1, 1, 1], [1, 1, 1]]),
+        ([1], [3, 3], None, 'numpy'),
+
+        # shape broadcasting
+        ([1], [1, 2], [0], 'explicit'),
+        ([1], [1, 2], [-2], 'explicit'),
+        ([1, 7], [5, 1, 7, 3], [1, 2], 'explicit'),
+        ([2, 1, 3], [2, 1, 3, 3], [0, 1, 2], 'explicit'),
+        ([2, 1, 3], [5, 2, 1, 3], [1, 2, 3], 'explicit'),
+
+        # value broadcasting
+        ([1], [1, 2], [0], 'explicit', [[1, 1]]),
+
+        ([[3, 1]], [2, 1, 2], [1, 2], 'explicit', [[[3, 1]], [[3, 1]]]),  # ref_shape (2, 1, 2)
+
+        ([[3, 1]], [2, 1, 2], [-2, -1], 'explicit', [[[3, 1]], [[3, 1]]]),  # ref_shape (2, 1, 2)
+
+        ([[[9, 5, 7]], [[9, 5, 7]]], [2, 2, 1, 3], [1, 2, 3], 'explicit',  # in_shape (2, 1, 3)
+         [[[[9, 5, 7]], [[9, 5, 7]]], [[[9, 5, 7]], [[9, 5, 7]]]]),        # ref_out_shape (2, 2, 1, 3)
+
+        ([[[9, 5, 7]], [[3, 4, 8]]], [2, 1, 3, 3], [0, 1, 2], 'explicit',             # in_shape (2, 1, 3)
+         [[[[9, 9, 9], [5, 5, 5], [7, 7, 7]]], [[[3, 3, 3], [4, 4, 4], [8, 8, 8]]]]), # ref_out_shape (2, 1, 3, 3)
+
+        # negative tests
+        ([1], [2, 2], [0], 'explicit', None, True),
+        ([1, 7], [5, 2, 7, 3], [1, 2], 'explicit', None, True),
+        ([1, 7], [5, 2, 7, 3], [2, 1], 'explicit', None, True),
+        ([1, 7], [5, 2, 7, 3], [-3, -2], 'explicit', None, True),
+    ])
+    def test_broadcast(self, data, target_shape, axes_mapping=None, mode='numpy', ref_out=None, test_raising=False):
+        if ref_out is not None:
+            input = valued_const_with_data('data', int64_array(data))
+        else:
+            input = shaped_data('data', int64_array(data))
+
+        nodes = {
+            **input,
+            **valued_const_with_data('target_shape', int64_array(target_shape)),
+            **regular_op_with_empty_data('broadcast', {'op': 'Broadcast', 'mode': mode}),
+        }
+
+        edges = [('data', 'broadcast'),
+                 ('target_shape', 'broadcast'),
+                 ('broadcast', 'broadcast_d')]
+
+        if axes_mapping is not None:
+            nodes.update(**valued_const_with_data('axes_mapping', int64_array(axes_mapping)))
+            edges.append(('axes_mapping', 'broadcast'))
+        graph = build_graph(nodes, edges)
+
+        broadcast_node = Node(graph, 'broadcast')
+        if test_raising:
+            self.assertRaises(AssertionError, Broadcast.infer, broadcast_node)
+            return
+
+        Broadcast.infer(broadcast_node)
+        if ref_out is not None:
+            self.assertTrue(np.array_equal(broadcast_node.out_node().value, np.array(ref_out)))
+        else:
+            self.assertTrue(np.array_equal(broadcast_node.out_node().shape, np.array(target_shape)))

model-optimizer/mo/utils/broadcasting.py

@@ -79,6 +79,26 @@ def bi_directional_shape_broadcasting(input_shape_1: np.array, input_shape_2: np
     return np.maximum(shape_1, shape_2)
 
 
+def explicit_shape_broadcasting(input_shape: np.array, target_shape: np.array, axes_mapping: np.array) -> np.array:
+    """
+    Explicit shape broadcasting of input tensor. Function only asserts that values are correct and normalizes axes.
+    Resulting shape is equal to target_shape.
+    :param input_value: input value to broadcast
+    :param target_shape: target shape
+    :param axes_mapping: a list of axis indices, each index maps an axis from the input_value to axis in the output
+    :return: broadcasted shape and normalized axes
+    """
+    assert np.all(np.diff(axes_mapping) >= 0), "axes_mapping is not sorted"
+    assert len(axes_mapping) == len(input_shape), "size of axes_mapping does not match to rank of input"
+    axes_mapping = np.array(list(map(lambda axis: axis + len(target_shape) if axis < 0 else axis, axes_mapping)))
+
+    res = target_shape.copy()
+    for i, axis in enumerate(axes_mapping):
+        assert 0 <= axis < len(res), "axis value from axes_mapping exceeds rank of target_shape"
+        assert res[axis] == input_shape[i], "specified mapping axis in target_shape differs from axis in input_shape"
+    return res, axes_mapping
+
+
 def uni_directional_broadcasting(input_value: np.array, target_shape: np.array):
     """
     Uni-directional broadcasting of input tensor to target shape following the numpy semantic
@@ -103,3 +123,18 @@ def bi_directional_broadcasting(input_value: np.array, second_shape: np.array):
         'The tensor of shape "{}" cannot be bi-directionally broadcasted to shape "{}"'.format(input_value.shape,
                                                                                                second_shape)
     return np.array(input_value * np.ones(second_shape), dtype=input_value.dtype)
+
+
+def explicit_broadcasting(input_value: np.array, target_shape: np.array, axes_mapping: np.array) -> np.array:
+    """
+    Explicit broadcasting of input tensor. Resulting shape is equal to target_shape except for axes specified in axes_mapping
+    :param input_value: input value to broadcast
+    :param target_shape: target shape
+    :param axes_mapping: a list of axis indices, each index maps an axis from the input_value to axis in the output
+    :return: broadcasted value
+    """
+    res_shape, normalized_axes_mapping = explicit_shape_broadcasting(input_value.shape, target_shape, axes_mapping)
+    expand_dim_axis = set(np.arange(len(target_shape))) - set(normalized_axes_mapping)
+
+    input_expanded = np.expand_dims(input_value.copy(), axis=list(expand_dim_axis))
+    return np.broadcast_to(input_expanded, res_shape)