def ti_set_output_port_shape(cycle_node, internal_id, port_num, iterations_count, axis): int_node_name = TensorIterator.find_internal_layer_id(cycle_node.body, internal_id) int_node = Node(cycle_node.body, int_node_name) assert int_node.op == 'Result' out_shape = int_node.in_port(0).data.get_shape().copy() # inside cycle node Unsqueeze was added to have the first dimension for concatenating results along it assert len(out_shape) >= 1 if axis is not None: out_shape[axis] = iterations_count assert port_num in cycle_node.out_ports() cycle_node.out_port(port_num).data.set_shape(out_shape)
def generate_port_map(node: Node, src_port_map, dir: str): result_list = [] for record in src_port_map: # do not update ids for not-connected output which is used in the Loop operation only if record['external_port_id'] != -1: if dir == 'out': # increase the output port id by the number of input ports # update the port id for proper generation of a "ports" section record['external_port_id'] += len(node.in_ports()) record[ 'internal_layer_id'] = TensorIterator.find_internal_layer_id( node.body, record['internal_layer_id']) result_list.append(record) return result_list