def test_splitv_zero(self):
graph = build_graph(self.nodes, self.edges,
{
'split_input_data': {'shape': int64_array([2, 12, 25, 30])},
'split_op': {'axis': np.array(2), 'split_lengths': np.array([2, 13, 10, 0]),
'out_ports_count': 4},
}
)
node = Node(graph, 'split_op')
for p in range(len(node.out_edges()), node.out_ports_count):
node.add_output_port(p)
AttributedVariadicSplit.infer(node)
self.assertTrue(len(node.out_edges()) == 3)
self.assertTrue(np.all(node.split_lengths == np.array([2, 13, 10])))
def test_splitv_zero_not_last(self):
graph = build_graph(self.nodes, self.edges,
{
'split_input_data': {'shape': int64_array([2, 12, 25, 30])},
'split_op': {'axis': np.array(2), 'split_lengths': np.array([2, 13, 0, 10]),
'out_ports_count': 4},
}
)
node = Node(graph, 'split_op')
# extractor should do it
for p in range(len(node.out_edges()), node.out_ports_count):
node.add_output_port(p)
node.out_port(2).get_connection().set_source(node.out_port(3))
AttributedVariadicSplit.infer(node)
self.assertTrue(node.out_port(3).disconnected())
self.assertTrue(np.all(node.split_lengths == np.array([2, 13, 10])))