enum class IrExportStage {
  HLO,
  HLO_NO_METADATA,
  HLO_SERIALIZED,
  OPTIMIZED_HLO,
  OPTIMIZED_HLO_SERIALIZED,
  OPTIMIZED_HLO_PROTO_SERIALIZED,
  OPTIMIZED_HLO_DOT
};

struct ArgShapeAndDType {
  TensorShape shape;
  DataType dtype;
};

enum class CompilerArgSource {
  TENSOR_SPEC,
  CONCRETE_INPUT,
};

// Returns the IR format of the selected stage for a given function `func_name`

BuildXlaCompilerArgumentFromTensorSpec(
    const FunctionBody* fbody, absl::Span<int const> must_be_constant_idxs,
    absl::Span<const Tensor* const> inputs,
    absl::Span<VariableInfo const> variable_args,
    absl::Span<const ArgShapeAndDType> flat_arg_shape_and_dtype) {
  TF_RET_CHECK(fbody != nullptr);
  auto& input_args = fbody->record->fdef().signature().input_arg();
  int input_arg_size = input_args.size();
  std::vector<XlaCompiler::Argument> args;