def test_topk_infer_opset1(self):
topk_node = Node(self.graph, 'topk')
topk_node['version'] = 'opset1'
TopK.infer(topk_node)
TopK.type_infer(topk_node)
self.assertTrue(np.array_equal(topk_node.out_port(0).data.get_shape(), int64_array([20, 10, 4])))
self.assertTrue(np.array_equal(topk_node.out_port(1).data.get_shape(), int64_array([20, 10, 4])))
self.assertTrue(topk_node.out_port(0).get_data_type() == np.float32)
self.assertTrue(topk_node.out_port(1).get_data_type() == np.int32)
def test_topk_infer_v10_i32_opset3(self):
self.graph.graph['cmd_params'] = FakeAttr(ir_version=10)
topk_node = Node(self.graph, 'topk')
topk_node['version'] = 'opset3'
topk_node['index_element_type'] = np.int32
TopK.infer(topk_node)
TopK.type_infer(topk_node)
self.assertTrue(np.array_equal(topk_node.out_port(0).data.get_shape(), int64_array([20, 10, 4])))
self.assertTrue(np.array_equal(topk_node.out_port(1).data.get_shape(), int64_array([20, 10, 4])))
self.assertTrue(topk_node.out_port(0).get_data_type() == np.float32)
self.assertTrue(topk_node.out_port(1).get_data_type() == np.int32)
def test_topk_infer_v10_opset1(self):
self.graph.graph['cmd_params'] = FakeAttr(
generate_experimental_IR_V10=True, ir_version=10)
topk_node = Node(self.graph, 'topk')
topk_node['version'] = 'opset1'
TopK.infer(topk_node)
TopK.type_infer(topk_node)
self.assertTrue(
np.array_equal(
topk_node.out_port(0).data.get_shape(),
int64_array([20, 10, 4])))
self.assertTrue(
np.array_equal(
topk_node.out_port(1).data.get_shape(),
int64_array([20, 10, 4])))
self.assertTrue(topk_node.out_port(0).get_data_type() == np.float32)
self.assertTrue(topk_node.out_port(1).get_data_type() == np.int32)