Fix ScatterND validation and implement evaluate (#4905)

* Fix ScatterND validation and implement evaluate

* Apply review feedback

* Update scatternd.py

model-optimizer/extensions/ops/scatternd.py

@@ -44,14 +44,16 @@ class ScatterNDBase(Op):
         # 1. ranks of both input and indices must be at least 1
         assert len(input_shape) >= 1 and len(indices_shape) >= 1, \
             'The node "{}" input and indices ranks must be at least 1'.format(node_name)
-           
+
         # 2. the last dimension of indices shape must be at most a rank of input
         assert indices_shape[-1] <= len(input_shape), \
             'The last dimension of indices shape must be at most a rank of input for the node "{}"'.format(node_name)
 
         # 3. updates is a tensor of shape indices_shape[:-1] + input_shape[indices_shape[-1]:]
+        # if expected updates shape is scalar, updates can be tensor with the single element (for example, of shape [1], [[1]], etc.)
         expected_updates_shape = np.concatenate((indices_shape[:-1], input_shape[indices_shape[-1]:]), axis=0)
-        assert np.array_equal(updates_shape, expected_updates_shape), \
+        assert np.array_equal(updates_shape, expected_updates_shape) or\
+               np.array_equal(expected_updates_shape, []) and np.array_equal(updates_shape, np.ones(len(updates_shape))), \
             'The updates shape must be equal to indices_shape[:-1] + input_shape[indices_shape[-1]:] for the node "{}"'.format(node_name)
 
         node.out_port(0).data.set_shape(input_shape)

model-optimizer/extensions/ops/scatternd_test.py

@@ -62,6 +62,11 @@ inputs7 = {'input': {'shape': int64_array([8]), 'value': int64_array([1, 2, 3, 4
           'updates': {'shape': int64_array([]), 'value': 9}}
 output7 = int64_array([1, 2, 3, 4, 9, 6, 7, 8])
 
+inputs8 = {'input': {'shape': int64_array([3]), 'value': int64_array([1, 2, 3])},
+          'indices': {'shape': int64_array([1]), 'value': int64_array([2])},
+          'updates': {'shape': int64_array([1]), 'value': int64_array([9])}}
+output8 = int64_array([1, 2, 9])
+
 class TestScatterNDUpdate(unittest.TestCase):
     def test_partial_infer1(self):
         graph = build_graph(nodes_attributes, edges, inputs1)
@@ -139,7 +144,7 @@ class TestScatterNDUpdate(unittest.TestCase):
         res_output_value = graph.node['output']['value']
 
         self.assertTrue(np.array_equal(output6, res_output_value),
-                        'values do not match expected: {} and given: {}'.format(output5, res_output_value))
+                        'values do not match expected: {} and given: {}'.format(output6, res_output_value))
 
     def test_infer7_scalar(self):
         graph = build_graph(nodes_attributes, edges, inputs7)
@@ -150,4 +155,15 @@ class TestScatterNDUpdate(unittest.TestCase):
         res_output_value = graph.node['output']['value']
 
         self.assertTrue(np.array_equal(output7, res_output_value),
-                        'values do not match expected: {} and given: {}'.format(output5, res_output_value))
+                        'values do not match expected: {} and given: {}'.format(output7, res_output_value))
+
+    def test_infer8(self):
+        graph = build_graph(nodes_attributes, edges, inputs8)
+        scatternd_node = Node(graph, 'scatternd_node')
+        ScatterNDUpdate.infer(scatternd_node)
+
+        # get the result
+        res_output_value = graph.node['output']['value']
+
+        self.assertTrue(np.array_equal(output8, res_output_value),
+                        'values do not match expected: {} and given: {}'.format(output8, res_output_value))