def test_expand_dims_infer_one_input_negative(self):
graph = build_graph(nodes_attributes,
[('input_1', 'expand_dims'),
('expand_dims', 'out')],
{'input_1': {'shape': np.array([3, 256, 256])},
'expand_dims': {'expand_axis': None}
})
expand_dims_node = Node(graph, 'expand_dims')
tf_expand_dims_infer(expand_dims_node)
self.assertIsNone(expand_dims_node.out_node().shape)
def test_expand_dims_infer_two_inputs_negative(self):
graph = build_graph(nodes_attributes,
[('input_1', 'expand_dims'),
('input_2', 'expand_dims'),
('expand_dims', 'out')],
{'input_1': {'shape': np.array([3, 256, 256])},
'input_2': {'shape': np.array([1]), 'value': np.array([2, 3], dtype=np.int32)},
})
expand_dims_node = Node(graph, 'expand_dims')
tf_expand_dims_infer(expand_dims_node)
self.assertIsNone(expand_dims_node.out_node().shape)
def test_expand_dims_infer_one_input_2(self):
graph = build_graph(nodes_attributes,
[('input_1', 'expand_dims'),
('expand_dims', 'out')],
{'input_1': {'shape': np.array([3, 256, 256])},
'expand_dims': {'expand_axis': 2}
})
expand_dims_node = Node(graph, 'expand_dims')
tf_expand_dims_infer(expand_dims_node)
exp_shape = np.array([3, 256, 1, 256])
res_shape = expand_dims_node.out_node().shape
self.assertEqual(len(exp_shape), len(res_shape))
for i in range(0, len(exp_shape)):
self.assertEqual(exp_shape[i], res_shape[i])
def test_expand_dims_infer_two_inputs_3(self):
graph = build_graph(nodes_attributes,
[('input_1', 'expand_dims'),
('input_2', 'expand_dims'),
('expand_dims', 'out')],
{'input_1': {'shape': np.array([3, 256, 256])},
'input_2': {'shape': np.array([]), 'value': np.array(3, dtype=np.int32)},
})
expand_dims_node = Node(graph, 'expand_dims')
tf_expand_dims_infer(expand_dims_node)
exp_shape = np.array([3, 256, 256, 1])
res_shape = expand_dims_node.out_node().shape
self.assertEqual(len(exp_shape), len(res_shape))
for i in range(0, len(exp_shape)):
self.assertEqual(exp_shape[i], res_shape[i])
