auto metadata_arg = metadata_proto.add_args();
    metadata_arg->set_kind(tpu::TPUCompileMetadataProto::Arg::PARAMETER);
    metadata_arg->set_dtype(dtype);
  }

  return absl::OkStatus();
}

absl::Status SetupReturnValues(mlir::ModuleOp module,
                               tpu::TPUCompileMetadataProto& metadata_proto) {
  auto main_fn = module.lookupSymbol<mlir::func::FuncOp>(kEntryFuncName);
  if (!main_fn) {