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)
