std::unique_ptr<Graph> g(new Graph(OpRegistry::Global()));
    TF_CHECK_OK(s.ToGraph(g.get()));
    std::cout << "Graph is " << (*g).ToGraphDefDebug().DebugString()
              << std::endl;
    auto node_name_image = g->BuildNodeNameIndex();
    node_name_image["identity0"]->AddAttr("_oc", "0");
    node_name_image["identity1"]->AddAttr("_oc", "1");

  }
  if (VLOG_IS_ON(4)) {
    for (Operation* new_control_predecessor : new_control_predecessors) {
        VLOG(4) << "      Adding predecessor op "
                << mlir::debugString(*new_control_predecessor);
    }
  }
  // Add new control predecessors to map.
  control_predecessors_[op].insert(new_control_predecessors.begin(),
                                   new_control_predecessors.end());

                         sizeof(value), os);
  }
};

// We print a protobuf using its ShortDebugString() when the string
// doesn't exceed this many characters; otherwise we print it using
// DebugString() for better readability.
const size_t kProtobufOneLinerMaxLength = 50;

template <typename T>
class TypeWithoutFormatter<T, kProtobuf> {
 public:
  static void PrintValue(const T& value, ::std::ostream* os) {

                         sizeof(value), os);
  }
};

// We print a protobuf using its ShortDebugString() when the string
// doesn't exceed this many characters; otherwise we print it using
// DebugString() for better readability.
const size_t kProtobufOneLinerMaxLength = 50;

template <typename T>
class TypeWithoutFormatter<T, kProtobuf> {
 public:
  static void PrintValue(const T& value, ::std::ostream* os) {

  Tensor parsed(tensor_proto.dtype());
  if (!parsed.FromProto(cpu_allocator(), tensor_proto)) {
    return errors::InvalidArgument("Cannot parse tensor from proto: ",
                                   tensor_proto.DebugString());
  }

  Status status;
  if (alloc_attrs.on_host()) {
    *tensor = parsed;
  } else {
    Allocator* allocator;
    {
      mutex_lock lock(mu_);

    return;
  }
  if (!ic_dst->MergeInput(dst.index, shape)) {
    status->status = tensorflow::errors::InvalidArgument(
        "Cannot update edge, incompatible shapes: ", ic_dst->DebugString(shape),
        " and ", ic_dst->DebugString(ic_dst->input(dst.index)), ".");
    return;
  }
  status->status = graph->graph.UpdateEdge(&new_src.oper->node, new_src.index,

        }
      });

  for (auto& src_arg : src_arg_pairs) {
    if (guaranteed_const_nodes.count(src_arg.first) != 0) {
      VLOG(1) << "Guaranteed const found: " << src_arg.first->DebugString();
      src_arg.second->AddAttr("_is_guaranteed_constant", true);
    }
  }
}

struct OutputInputTensorPairHasher {
  uint64 operator()(std::pair<OutputTensor, InputTensor> const& s) const {

      TF_RETURN_IF_ERROR(ProcessWhileLoop(n, dy[0]));
      continue;
    }
    // All loop-specific control flow ops should have been handled above
    DCHECK(!n->IsEnter() && !n->IsNextIteration()) << n->DebugString();

    const int num_no_grad = no_grad_dy_indices.size();
    if (IsPrimitiveOpWithNoGrad(n->type_string()) || num_no_grad == num_y) {
      // No grad defined for this op, or all outputs returned 'NoGradient':

    }
  }

  if (dimension_to_splits_map.empty()) {
    return absl::InvalidArgumentError(absl::StrCat(
        "Arg has unnecessary tiled sharding: ", sharding.DebugString()));
  }

  return dimension_to_splits_map;
}

int GetDimsFromXLAShardingTiled(const xla::OpSharding& xla_sharding) {
  return xla_sharding.tile_assignment_dimensions_size() -

            << "tensor of type " << ranked_type << " (rank=" << tensor_rank
            << ") sharded in " << (tile_assignment_rank - tensor_rank)
            << " extra dimension(s) by: " << sharding->DebugString();
      }

      return mlir::failure();
    }
  }
  return mlir::success();
}

// Verify sharding for all arguments and return values.
LogicalResult VerifyShardings(mlir::func::FuncOp func,