void Weaver::Finalize()

in tensorflow_fold/loom/weaver.cc [716:797]
53 lines of code
12 McCabe index (conditional complexity)

void Weaver::Finalize() {
  // Make sure Finalize only does anything the first time.
  if (finalized_) return;
  finalized_ = true;

  if (max_depth_ == -1) {
    max_depth_ = Deepest();
  }

  // Add PassThroughs so that all the outputs are at the 'max_depth_'.
  std::vector<tensor_idx_t> deepened_outputs;
  for (tensor_idx_t output_result_id : output_result_ids_) {
    deepened_outputs.push_back(Deepen(output_result_id, max_depth_));
  }

  // Given a depth, an op_idx and an op_output_idx, wiring_offset tells us where
  // in the state Tensor of the appropriate type_shape the 'op_output_idx'th
  // output of 'op' will start, after the concat which takes place in sub-layer
  // (3) of the previous loom layer.
  std::map<std::tuple<tensor_idx_t, tensor_idx_t, tensor_idx_t>,
           tensor_idx_t> wiring_offset;

  // Named tensors, constants and batch input members have depth=0, op_idx=-1
  // and op_output_idx=-1, and we don't have to add anything to their 'pos_idx'
  // to get their index in the initial state tensor.
  wiring_offset[std::make_tuple(0, -1, -1)] = 0;

  for (tensor_idx_t depth = 1; depth <= max_depth_; ++depth) {
    std::vector<tensor_idx_t> concat_offset(num_type_shapes_);
    for (tensor_idx_t op_idx = 0; op_idx < num_ops_; ++op_idx) {
      // Note: the number of times op 'op_idx' got called at 'depth' is the same
      // as the number of 0th arguments passed in.
      tensor_idx_t num_op_calls = wiring_results_[
            std::make_tuple(depth, op_idx, 0)].size();

      for (tensor_idx_t output_idx = 0;
           output_idx < op_output_ts_idx_[op_idx].size();
           ++output_idx) {
        // The TypeShape for the current output determines which state tensor
        // this batch of outputs ended up in.
        tensor_idx_t ts_idx = op_output_ts_idx_[op_idx][output_idx];

        wiring_offset[std::make_tuple(depth, op_idx, output_idx)] =
            concat_offset[ts_idx];

        // We're incrementing concat_offset by 'num_op_calls' because the
        // current op must have 'num_op_calls' as the batch size for all of its
        // outputs (including this one.)
        concat_offset[ts_idx] += num_op_calls;
      }
    }
  }

  // Populate 'final_wiring_' with the values that will end up in the wiring
  // diagram placeholder variables in the loom.
  for (tensor_idx_t depth = 1; depth <= max_depth_; ++depth) {
    for (tensor_idx_t op_idx = 0; op_idx < num_ops_; ++op_idx) {
      for (tensor_idx_t arg_idx = 0; arg_idx < op_input_ts_idx_[op_idx].size();
           ++arg_idx) {
        auto key = std::make_tuple(depth, op_idx, arg_idx);
        const std::vector<tensor_idx_t> &my_wiring_results =
            wiring_results_[key];
        std::vector<tensor_idx_t> &my_final_wiring = final_wiring_[key];
        for (tensor_idx_t result_id : my_wiring_results) {
          const LoomResult &r = loom_results_[result_id];
          my_final_wiring.push_back(
              r.pos_idx + wiring_offset[std::make_tuple(
                  r.depth, r.op_idx, r.op_output_idx)]);
        }
      }
    }
  }

  // Populate 'final_output_wiring_'
  final_output_wiring_.resize(num_type_shapes_);
  for (tensor_idx_t deepened_output : deepened_outputs) {
    const LoomResult &r = loom_results_[deepened_output];
    final_output_wiring_[r.ts_idx].push_back(
        r.pos_idx + wiring_offset[std::make_tuple(
            r.depth, r.op_idx, r.op_output_idx)]);
  }
}