def test_node_factory_topk(): dtype = np.int32 data = ng.parameter([2, 10], dtype=dtype, name="A") k = ng.constant(3, dtype=dtype, name="B") factory = _NodeFactory("opset1") node = factory.create("TopK", [data, k], {"axis": 1, "mode": "max", "sort": "value"}) assert node.get_type_name() == "TopK" assert node.get_output_size() == 2 assert list(node.get_output_shape(0)) == [2, 3]
def test_node_factory_add(): shape = [2, 2] dtype = np.int8 parameter_a = ng.parameter(shape, dtype=dtype, name="A") parameter_b = ng.parameter(shape, dtype=dtype, name="B") factory = _NodeFactory("opset1") node = factory.create("Add", [parameter_a, parameter_b], {}) assert node.get_type_name() == "Add" assert node.get_output_size() == 1 assert list(node.get_output_shape(0)) == [2, 2]
def test_node_factory_add(): shape = [2, 2] dtype = np.int8 parameter_a = ng.parameter(shape, dtype=dtype, name='A') parameter_b = ng.parameter(shape, dtype=dtype, name='B') factory = _NodeFactory('opset1') node = factory.create('Add', [parameter_a, parameter_b], {}) assert node.get_type_name() == 'Add' assert node.get_output_size() == 1 assert list(node.get_output_shape(0)) == [2, 2]
def test_node_factory_topk(): dtype = np.int32 data = ng.parameter([2, 10], dtype=dtype, name='A') k = ng.constant(3, dtype=dtype, name='B') factory = _NodeFactory('opset1') node = factory.create('TopK', [data, k], { 'axis': 1, 'mode': 'max', 'sort': 'value' }) assert node.get_type_name() == 'TopK' assert node.get_output_size() == 2 assert list(node.get_output_shape(0)) == [2, 3]
def test_node_factory_topk(): dtype = np.int32 data = ng.parameter([2, 10], dtype=dtype, name="A") k = ng.constant(3, dtype=dtype, name="B") factory = _NodeFactory("opset1") arguments = NodeFactory._arguments_as_outputs([data, k]) node = factory.create( "TopK", arguments, {"axis": 1, "mode": "max", "sort": "value"} ) attributes = node.get_attributes() assert node.get_type_name() == "TopK" assert node.get_output_size() == 2 assert list(node.get_output_shape(0)) == [2, 3] assert attributes["axis"] == 1 assert attributes["mode"] == "max" assert attributes["sort"] == "value"
def __init__(self, opset_version: str = DEFAULT_OPSET) -> None: """Create the NodeFactory object. @param opset_version: The opset version the factory will use to produce ops from. """ self.factory = _NodeFactory(opset_version)
def __init__(self, opset_version='opset1'): # type: (str) -> None """Create the NodeFactory object. :param opset_version: The opset version the factory will use to produce ops from. """ self.factory = _NodeFactory(opset_version)