merge_node.cc

// Original : https://github.com/alpa-projects/tensorflow-alpa/blob/d298f84474a04ecce02085332793e6115c0c8e0e/tensorflow/compiler/xla/service/spmd/auto_sharding_strategy.h#L854-L876

    if (adj_list.size() > 1) {
      //  Merge src to dst.
      //
      //  Before:
      //
      //  src ---- adj ---- dst
      //   |                 |
      //   -------------------
      //
      //  After:
      //
      //           adj ---- dst
      //
      for (int adj : adj_list) {
        if (adj == dst) {
          continue;
        }
        Matrix added_edge_cost(node_lens[dst], node_lens[adj]);

        for (int i = 0; i < node_lens[dst]; ++i) {
          int j = reindexing[i];
          Matrix edge_cost_src_adj = GetEdgeCost(src, adj);
          for (int k = 0; k < node_lens[adj]; ++k) {
            added_edge_cost(i, k) = edge_cost(i, j) + edge_cost_src_adj(j, k);
          }
        }

        AddEdgeCost(dst, adj, added_edge_cost);
      }
    } else {
      //  Merge src to dst.
      //
      //  Before:
      //
      //  src ---- dst
      //
      //  After:
      //
      //           dst
      //
      for (int i = 0; i < node_lens[dst]; ++i) {
        int j = reindexing[i];
        extra_node_costs[dst][i] += edge_cost(i, j);
      }
    }



// Updated:
    for (int adj : adj_list) {
      if (adj == dst) {
        for (int i = 0; i < node_lens[dst]; ++i) {
          int j = reindexing[i];
          extra_node_costs[dst][i] += edge_cost(i, j);
        }
      } else {
        Matrix added_edge_cost(node_lens[dst], node_lens[adj]);

        for (int i = 0; i < node_lens[dst]; ++i) {
          int j = reindexing[i];
          Matrix edge_cost_src_adj = GetEdgeCost(src, adj);
          for (int k = 0; k < node_lens[adj]; ++k) {
            added_edge_cost(i, k) = edge_cost_src_adj(j, k);
          }
        }

        AddEdgeCost(dst, adj, added_edge_cost);
      }
    }