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`