def test_tile_infer_all_ones(self):
graph = build_graph(nodes_attributes, edges,
{'tile_values': {
'value': np.array([1, 1, 1, 1])
}})
tile_node = Node(graph, 'tile')
Tile.infer(tile_node)
self.assertTrue(
np.all(
np.array([10, 20, 30, 40]) == graph.node['tile_out']['shape']))
def test_tile_infer_three_non_one(self):
graph = build_graph(nodes_attributes, edges,
{'tile_values': {
'value': np.array([2, 1, 5, 2])
}})
tile_node = Node(graph, 'tile')
Tile.infer(tile_node)
self.assertTrue(
np.all(
np.array([20, 20, 150, 80]) == graph.node['tile_out']
['shape']))
def test_tile_infer_shapes_alignment(self):
graph = build_graph(nodes_attributes, edges, {
'tile_values': {
'value': np.array([1, 2, 3]),
'shape': np.array([3])
}
})
tile_node = Node(graph, 'tile')
Tile.infer(tile_node)
self.assertTrue(
np.all(
np.array([10, 20, 60, 120]) == graph.node['tile_out']
['shape']))
def test_tile_infer_correct_2d_tensor(self):
graph = build_graph(
nodes_attributes, edges, {
'data': {
'shape': np.array([3, 7])
},
'tile_values': {
'value': np.array([5, 1])
}
})
tile_node = Node(graph, 'tile')
Tile.infer(tile_node)
self.assertTrue(
np.all(np.array([15, 7]) == graph.node['tile_out']['shape']))
def test_tile_infer_values_test(self):
input_data = np.arange(-30, 60, 0.25).reshape([2, 4, 3, -1])
tile_values = np.array([3, 1, 1, 1])
graph = build_graph(
nodes_attributes, edges, {
'data': {
'shape': np.array(input_data.shape),
'value': input_data
},
'tile_values': {
'value': tile_values
}
})
tile_node = Node(graph, 'tile')
Tile.infer(tile_node)
self.assertTrue(
np.all(
np.tile(input_data, tile_values) == graph.node['tile_out']
['value']))
def replace_pattern(self, graph: Graph, match: dict):
node = match['tile']
name = node.soft_get('name', node.id)
axis = node.axis
tiles = node.tiles
input_shape = node.in_port(0).data.get_shape()
assert input_shape is not None
tiles_input_value = int64_array(np.ones(input_shape.size))
tiles_input_value[axis] = tiles
const = Const(graph, {
'value': tiles_input_value,
'name': name + '/tiles'
}).create_node()
tile = Tile(graph, {'name': name}).create_node()
node.out_port(0).get_connection().set_source(tile.out_port(0))
node.in_port(0).get_connection().set_destination(tile.in_port(0))
const.out_port(0).connect(tile.in_port(1))
def test_tile_infer_values_const_propagation(self):
"""
Test for constant propagation even if tile with multiple tile indices is not supported
"""
input_data = np.arange(-30, 60, 0.25).reshape([2, 4, 3, -1])
tile_values = np.array([4, 3, 2, 5])
graph = build_graph(
nodes_attributes, edges, {
'data': {
'shape': np.array(input_data.shape),
'value': input_data
},
'tile_values': {
'value': tile_values
}
})
tile_node = Node(graph, 'tile')
Tile.infer(tile_node)
self.assertTrue(
np.all(
np.tile(input_data, tile_values) == graph.node['tile_out']
['value']))
def extract(cls, node):
Tile.update_node_stat(node, {})
return cls.enabled
def mxrepeat_decomposition(node: Node):
graph = node.graph
name = node.soft_get('name', node.id)
rename_node(node, name + '/to_be_removed')
# Unqueeze
input_rank = Rank(graph, {'name': name + '/Rank'}).create_node()
node.in_port(0).get_source().connect(input_rank.in_port(0))
axis = get_canonical_axis_index_node(input_rank, node.axis)
unsqueeze_axis = create_op_node_with_second_input(
graph,
Add,
int64_array([1]), {'name': name + '/Unsqueeze/Axis'},
input_node=axis)
unsqueeze = Unsqueeze(graph, {
'name': name + '/Unsqueeze'
}).create_node()
unsqueeze.in_port(1).connect(unsqueeze_axis.out_port(0))
# Tile (1, 1, ..., repeats, ..., 1)
# we generate tile array according to the following table:
# parts: | first | repeats | second |
# i: | 0, 1, ..., axis,| axis + 1,| ..., rank+1 |
# tile_array: | 1, 1, ..., 1 ,| repeats ,| ..., 1 |
one = Const(graph, {
'name': name + '/Broadcast/One',
'value': int64_array([1])
}).create_node()
first_ones = Broadcast(graph, {
'name': name + '/Broadcast/Ones_first_part'
}).create_node()
first_ones.in_port(0).connect(one.out_port(0))
first_ones.in_port(1).connect(unsqueeze_axis.out_port(0))
repeats = Const(graph, {
'name': name + '/repeats',
'value': int64_array([node.repeats])
}).create_node()
second_ones = Broadcast(graph, {
'name': name + '/Broadcast/Ones_second_part'
}).create_node()
second_part_broadcast_shape = Sub(
graph, {
'name': name + '/Broadcast/Shape/second_part'
}).create_node()
second_part_broadcast_shape.in_port(0).connect(input_rank.out_port(0))
second_part_broadcast_shape.in_port(1).connect(
unsqueeze_axis.out_port(0))
second_ones.in_port(0).connect(one.out_port(0))
second_ones.in_port(1).connect(second_part_broadcast_shape.out_port(0))
tile_repeats = new_shape_node_from_shape_nodes(
[first_ones, repeats, second_ones])
tile = Tile(graph, {'name': name + '/Tile'}).create_node()
tile.in_port(1).connect(tile_repeats.out_port(0))
# Reshape (input_shape[:axis], input_shape[axis] * repeats, input_shape[axis+1:])
# we generate reshape dim array according to the following table:
# parts: | first | rep | second |
# i: | 0, 1, ... ,| axis, | ..., rank |
# dim_array: | inp_sh[i] ,| input_shape[axis] * repeats ,| inp_sh[i] |
input_shape = Shape(graph, {'name': name + '/Shape'}).create_node()
node.in_port(0).get_source().connect(input_shape.in_port(0))
first_input_shape_part = get_shape_values_by_range_idxs(
input_shape,
input_rank,
begin=0,
end=node.axis,
include_begin=True,
include_end=False)
original_axis_dim = create_op_with_const_inputs(
graph,
Gather, {2: int64_array(0)}, {'name': name + '/OriginalDim'},
input_node=input_shape)
original_axis_dim.in_port(1).connect(axis.out_port(0))
repeated_dimention = Mul(graph, {
'name': name + '/RepeatedDim'
}).create_node()
repeated_dimention.in_port(0).connect(original_axis_dim.out_port(0))
repeated_dimention.in_port(1).connect(repeats.out_port(0))
second_input_shape_part = get_shape_values_by_range_idxs(
input_shape,
input_rank,
begin=node.axis,
end=-1,
include_begin=False,
include_end=True)
output_shape = new_shape_node_from_shape_nodes([
first_input_shape_part, repeated_dimention, second_input_shape_part
])
reshape = Reshape(graph, {'name': name}).create_node()
rename_node(reshape, name)
reshape.in_port(1).connect(output_shape.out_port(0))
# Final connections
node.in_port(0).get_connection().set_destination(unsqueeze.in_port(0))
tile.in_port(0).connect(unsqueeze.out_port(0))
reshape.in_port(0).connect(tile.out_port(0))
node.out_port(0).get_connection().set_source(reshape.out_port(0))
def extract(cls, node: Node):
attrs = get_mxnet_layer_attrs(node.symbol_dict)
Tile.update_node_stat(node, {
'reps': attrs.tuple('reps', int, None),
})
return cls.enabled