def find_and_replace_pattern(self, graph: Graph): for ti in graph.get_op_nodes(type='TensorIterator'): self.external_nodes_normalization(ti) if len([record for record in ti.input_port_map if record.get('axis') is not None]) == 0: for record in ti.output_port_map: if record.get('axis') is not None: record['start'] = 0 real_output_port = TensorIterator.special_port_to_real_port(ti, record['external_port_id'], 'out') output_shape = ti.out_port(real_output_port).data.get_shape() assert output_shape is not None record['end'] = output_shape[record['axis']]
def external_nodes_normalization(ti): """ TensorIterator external ports may have several internal layer connections. Current transformation does the following: - normalizes port maps (eliminating duplicated records) - replicates external input/output port for each internal Parameter/Result it is connected to - updates input and output port maps according to previous step replications """ def update_external_port_id(ti, port_type, old_external_port_id, new_external_port_id, internal_layer_id): assert port_type in ['in', 'out'] port_map = ti.input_port_map if port_type == 'in' else ti.output_port_map for record in port_map: if record['external_port_id'] == old_external_port_id and \ record['internal_layer_id'] == internal_layer_id: record['external_port_id'] = new_external_port_id NormalizeTI.maps_uniqueization(ti) body = ti.body external_input_ports = defaultdict(list) for record in ti.input_port_map: assert 'external_port_id' in record external_input_ports[record['external_port_id']].append(record) for external_port_id, record_list in external_input_ports.items(): if len(record_list) == 1: continue real_external_port_id = TensorIterator.special_port_to_real_port( ti, external_port_id, 'in') source = ti.in_port(real_external_port_id).get_source() for record in record_list[1:]: assert 'internal_layer_id' in record new_real_input_port_id = max(map(int, ti.in_ports().keys())) + 1 new_external_port_id = max([ int(d['external_port_id']) for d in list(ti.in_edges().values()) + list(ti.out_edges().values()) ]) + 1 ti.add_input_port(new_real_input_port_id) source.connect(ti.in_port(new_real_input_port_id)) ti.in_edge(new_real_input_port_id )['external_port_id'] = new_external_port_id update_external_port_id(ti, 'in', external_port_id, new_external_port_id, record['internal_layer_id']) external_output_ports = defaultdict(list) for record in ti.output_port_map: assert 'external_port_id' in record external_output_ports[record['external_port_id']].append(record) for external_port_id, record_list in external_output_ports.items(): if len(record_list) == 1: continue real_external_port_id = TensorIterator.special_port_to_real_port( ti, external_port_id, 'out') dsts = ti.out_port(real_external_port_id).get_destinations() for record in record_list[1:]: assert 'internal_layer_id' in record new_real_output_port_id = max(map(int, ti.out_ports().keys())) + 1 new_external_port_id = max([ int(d['external_port_id']) for d in list(ti.in_edges().values()) + list(ti.out_edges().values()) ]) + 1 ti.add_output_port(new_real_output_port_id) for dst in dsts: ti.out_port(new_real_output_port_id).connect(dst) update_external_port_id(ti, 'out', external_port_id, new_external_port_id, record['internal_layer_id']) body.clean_up()