Function BuildComputation

tensorflow/compiler/tf2xla/xla_compiler.cc:155–363

Builds the XLA computation. - `args` is the list of input arguments - `retvals` is the list of retvals produced by _Retval operators, in index order. - `arg_shardings` and `retval_shardings` are mapping from arg/return indices to sharding. - If `return_updated_values_for_all_resources` is true, all resources will be included in `resource_updates`, regardless of whether their value changed. - Sets

//   resource variable argument to the computation to be updated, and `type` is
//   the type of the final output.
Status BuildComputation(
    const std::vector<XlaCompiler::Argument>& args,
    const std::vector<XlaExpression>& retvals,
    const std::map<int, xla::OpSharding>& arg_shardings,
    const std::map<int, xla::OpSharding>& retval_shardings,
    const std::vector<std::unique_ptr<XlaResource>>& resources,
    std::unique_ptr<xla::XlaOp> token_output,
    const XlaCompiler::ShapeRepresentationFn& shape_representation_fn,
    bool return_updated_values_for_all_resources, bool always_return_tuple,
    xla::XlaBuilder* builder, xla::XlaComputation* computation,
    int* num_computation_outputs, int* num_nonconst_outputs,
    std::vector<XlaCompiler::OutputDescription>* outputs,
    std::vector<XlaCompiler::ResourceUpdate>* resource_updates,
    xla::Shape* output_shape) {
  // Attach a common operator name as metadata. This has no semantic effect — it
  // merely makes the HLO graph more readable when visualized via TensorBoard,
  // since TensorBoard forms groups out of operators with similar names.
  xla::OpMetadata retval_metadata;
  retval_metadata.set_op_name("XLA_Retvals");
  builder->SetOpMetadata(retval_metadata);
  VLOG(1) << "Building new computation";
  auto cleanup = gtl::MakeCleanup([builder]() { builder->ClearOpMetadata(); });

  // Builds a no-op XLA computation. We need to set the sharding of outputs, but
  // cannot change the sharding of the existing output op. To do this, we build
  // a new identity op to which shardings can be applied.
  auto identity_op = [builder](xla::XlaOp op) {
    return xla::GetTupleElement(xla::Tuple(builder, {op}), 0);
  };

  std::vector<xla::XlaOp> elems;
  elems.reserve(retvals.size());

  // Keeps track of the layout of each retval. If a retval is not in this list,
  // a descending layout is used. The first element is the output index, second
  // element is the new layout.
  std::vector<std::pair<int64, xla::Layout>> retval_index_and_layout;
  for (int i = 0; i < retvals.size(); ++i) {
    XlaCompiler::OutputDescription& output = (*outputs)[i];
    const XlaExpression& retval = retvals[i];
    output.type = retval.dtype();
    switch (retval.kind()) {
      case XlaExpression::Kind::kConstant:
        output.is_constant = true;
        output.constant_value = retval.constant_value();
        output.shape = output.constant_value.shape();
        break;

      case XlaExpression::Kind::kTensorList: {
        output.is_tensor_list = true;
        xla::XlaOp value = retval.handle();
        elems.push_back(value);
        break;
      }

      case XlaExpression::Kind::kXlaOp: {
        output.is_constant = false;
        TF_ASSIGN_OR_RETURN(output.shape, retval.GetShape());