se::Stream* const src_compute_stream = src_xla_context->stream();
    TF_RET_CHECK(src_compute_stream != nullptr);
    TF_RET_CHECK(input->dtype() == output->dtype())
        << "input type: " << DataTypeString(input->dtype()) << " output type "
        << DataTypeString(output->dtype());
    TF_RET_CHECK(input->shape() == output->shape());
    TF_RET_CHECK(DMAHelper::CanUseDMA(input));
    auto* const src_compute_stream_impl =

    TF_RETURN_IF_ERROR(
        ctx->allocate_temp(output_dtype, output_shape, &output_tensor));
    if (output_tensor.TotalBytes() > 0) {
      XlaTensor* xla_tensor = XlaTensor::FromTensor(&output_tensor);
      TF_RET_CHECK(xla_tensor);
      xla_tensor->set_shaped_buffer(output.TakeSubTree({output_num}));
      if (use_multiple_streams) {
        xla_tensor->ResetDefinitionEvent(definition_event, stream);
      }
    }

  auto dict_attr =
      function->getAttrOfType<mlir::DictionaryAttr>(kEntryFuncAttr);
  if (dict_attr) {
    TF_RET_CHECK(mlir::isa<mlir::StringAttr>(dict_attr.get("inputs")))
        << "inputs missing in entry function attribute";
    TF_RET_CHECK(mlir::isa<mlir::StringAttr>(dict_attr.get("outputs")))
        << "outputs missing in entry function attribute";

  auto dict_attr =
      function->getAttrOfType<mlir::DictionaryAttr>(kEntryFuncAttr);
  if (dict_attr) {
    TF_RET_CHECK(mlir::isa<mlir::StringAttr>(dict_attr.get("inputs")))
        << "inputs missing in entry function attribute";
    TF_RET_CHECK(mlir::isa<mlir::StringAttr>(dict_attr.get("outputs")))
        << "outputs missing in entry function attribute";

    // Treats empty input shape as scalar
    if (dim_str.empty()) continue;
    if (dim_str == "?") {
      dims.push_back(-1);
      continue;
    }
    int size;
    TF_RET_CHECK(absl::SimpleAtoi(dim_str, &size));
    dims.push_back(size);
  }
  return dims;
}

static Status HandleSubtype(absl::string_view subtype,
                            ArrayInfo::SubTypeInfo* result) {

    // tf_library genrule are deterministic.
    const size_t size = module->ByteSizeLong();
    auto serialized = std::make_unique<char[]>(size);
    TF_RET_CHECK(
        SerializeToBufferDeterministic(*module, serialized.get(), size));
    TF_RETURN_IF_ERROR(
        WriteStringToFile(Env::Default(), flags.out_session_module,

                            device_tensor->shape(), device_tensor->dtype(),
                            /*fast_mem=*/false, layout_preference));

    // The device tensor should always be fresh.
    TF_RET_CHECK(!xla_tensor->has_shaped_buffer());

    TF_RETURN_IF_ERROR(
        xla_tensor->AllocateShapedBuffer(device_tensor->dtype(), shape, client_,
                                         stream_->parent()->device_ordinal()));

  xla::HloModuleProto* module_proto =
      compilation_result->computation->mutable_proto();
  absl::StatusOr<xla::ProgramShape> program_shape_or_status =
      compilation_result->computation->GetProgramShape();
  TF_RET_CHECK(program_shape_or_status.ok());

  xla::ProgramShape& program_shape = program_shape_or_status.value();
  if (!program_shape.result().IsTuple())
    return errors::Internal("Expect result to have tuple shape");

    // the attribute from MLIR, it is treated as an attribute from function
    // calls.
    std::vector<string> name_tokens =
        absl::StrSplit(name, '.', absl::SkipEmpty());
    TF_RET_CHECK(name_tokens.size() <= 2);
    auto it = func_call_attrs.find(name_tokens[0]);
    if (it == func_call_attrs.end()) {
      (*values)[string(name)] = std::move(value);
    } else {

    std::vector<xla::Shape>* xla_input_shapes) {
  xla_input_shapes->clear();

  mlir::func::FuncOp main_func =
      module.lookupSymbol<mlir::func::FuncOp>("main");
  TF_RET_CHECK(main_func != nullptr) << "No main function found";
  mlir::FunctionType func_type = main_func.getFunctionType();

  int num_args = func_type.getNumInputs();
  xla_input_shapes->reserve(num_args);