if (!parent_op) return false;
  return HasOutsideCompilationAncestor(parent_op);
}

Operation* ApplyXlaHostTransferAttr(Operation* op, OpBuilder& builder) {
  op->setAttr("_xla_has_host_transfer", builder.getBoolAttr(true));
  return op;
}

// Creates a tf._XlaSendFromHost or tf._XlaSendFromHostV2 op. If device ordinal
// is present, a tf._XlaSendFromHostV2 op is created instead.

      {
        ::mlir::Value x_value = (*x.begin());
        mean_op_1 = rewriter.create<TF::MeanOp>(
            odsLoc, new_mean_type, x_value, reduce_dim_op,
            /*keep_dims=*/rewriter.getBoolAttr(false));
      }

      ::mlir::TF::SquaredDifferenceOp square_diff_op;
      {
        ::mlir::Value tblgen_value_0 = (*x.begin());

        /*then_branch=*/
        mlir::SymbolRefAttr::get(then_branch_op),
        /*else_branch=*/
        mlir::SymbolRefAttr::get(else_branch_op),
        /*is_stateless=*/rewriter.getBoolAttr(true));
    return success();
  }

 private:
  // When the input tensorlist's size is smaller than the requested size,
  // then branch is executed.

    // gives mathematically equivalent result.
    auto max_logits =
        rewriter.create<TF::MaxOp>(loc, logits, reduce_dim,
                                   /*keep_dims=*/rewriter.getBoolAttr(true));
    auto shifted_logits = rewriter.create<TF::SubOp>(loc, logits, max_logits);

    // Exponentiate the inputs.
    Value exp = rewriter.create<TF::ExpOp>(loc, shifted_logits);

  return vec;
}

void SetBasicBlockAttributes(OpBuilder& builder, Operation* op) {
  op->setAttr(kDevice, builder.getStringAttr(""));
  op->setAttr(kLower, builder.getBoolAttr(true));
}

std::vector<Value> ResultsAsVector(Operation* op, int begin, int num) {
  int end = begin + num;
  std::vector<Value> vec;
  vec.reserve(end - begin);

        builder.setInsertionPointAfter(last_store);
        auto cast = builder.create<TF::CastOp>(
            last_store.getLoc(), read_type, last_store.getValue(),
            /*Truncate=*/builder.getBoolAttr(false));
        read_variable_op.getValue().replaceAllUsesWith(cast);
      } else {
        read_variable_op.getValue().replaceAllUsesWith(last_store.getValue());
      }

      read_variable_op.erase();

        builder.getNamedAttr("then_branch", then_attr),
        builder.getNamedAttr("else_branch", else_attr),
        // TODO(b/139667752): Analyze statelessness correctly
        builder.getNamedAttr("is_stateless", builder.getBoolAttr(false))};
  }
  if (auto* opts = options.AsWhileOptions()) {
    uint32_t cond_idx = opts->cond_subgraph_index;
    if (cond_idx >= func_names.size()) {