EXPECT_TRUE(succeeded(parse_result));

  return cast<ModuleOp>(block.front());
}

TEST_F(TfSavedModelTest,
       GetInitializerFunctionReturnsNullWhenNoSessionInitializerOp) {
  constexpr StringRef kModuleOpStr =
      R"mlir(module attributes {tf_saved_model.semantics} {})mlir";

  Block block;
  ModuleOp module_op = ParseModuleOp(kModuleOpStr, block, ctx_);

    dict_value: {
    }
  }
}

)";

constexpr absl::string_view kZeroReturnOutputSignatureTextProto = R"(
none_value: {}
)";

constexpr absl::string_view kSingleReturnOutputSignatureTextProto = R"(
tensor_spec_value: {
  shape: {
    dim: {
        size: 1
    }
  }
  dtype: DT_FLOAT
}
)";

constexpr absl::string_view kThreeReturnOutputSignatureTextProto = R"(

// for debugging.
inline constexpr absl::string_view kTfQuantPtqPreCalibrationStepName =
    "tf_quant_ptq_pre_calibration";
inline constexpr absl::string_view kTfQuantPtqPostCalibrationStepName =
    "tf_quant_ptq_post_calibration";
inline constexpr absl::string_view kTfQuantQatStepName = "tf_quant_qat";
inline constexpr absl::string_view kTfQuantPtqDynamicRangeStepName =
    "tf_quant_ptq_dynamic_range";

using ::mlir::mhlo::test::GetMlirModuleFromString;
using ::tensorflow::monitoring::testing::CellReader;

constexpr char kNotDynamicFunctionName[] = "kNotDynamicFunction";
constexpr char kDynamicFunctionName[] = "kDynamicFunction";
static constexpr char kDynamismOpCounterStreamzName[] =
    "/tensorflow/core/tf2xla/api/v2/dynamism_op_counter";
static constexpr char kDynamismFunctionCounterStreamzName[] =
    "/tensorflow/core/tf2xla/api/v2/dynamism_function_counter";

// Device types.
inline constexpr llvm::StringRef kTpuDevice = "TPU";
// _xla_outside_compilation
inline constexpr llvm::StringRef kXlaOutsideCompilationAttr =
    "_xla_outside_compilation";
// device attr
inline constexpr llvm::StringRef kDeviceAttr = "device";
// Function attribute to signal that a function should be skipped from TPU
// island outlining. The attribute is set in

// The main function's name in the serialized stablehlo module embedded in
// XlaCallModule's `module` attribute.
constexpr llvm::StringRef kStablehloMainFunctionName = "main";

// After deserializing the stablehlo functions from XlaCallModule,
// this XlaCallModule attribute refers to the deserialized stablehlo main
// function.
constexpr llvm::StringRef kStablehloEntryFunctionAttrName = "_entry_function";

  static constexpr size_t kResultIndex = 0;

  // Shapes of the input arguments.
  static constexpr int32_t kArg0Shapes[] = {
1, 2
  };
  static constexpr int32_t kArg1Shapes[] = {
3, 4
  };
  static constexpr int32_t kArg2Shapes[] = {
1
  };
  static constexpr int32_t kArg3Shapes[] = {
1
  };
  static constexpr int32_t kArg4Shapes[] = {
5
  };

using tpu::FunctionToHloArgs;
using tpu::MlirToHloArgs;
using tpu::ShardingAndIndex;
using tpu::TPUCompileMetadataProto;

static constexpr char kCompilationTimeStreamzName[] =
    "/tensorflow/core/tf2xla/api/v2/phase2_compilation_time";
static constexpr char kFullBridge[] = "full_bridge";
static constexpr char kCompilationStatusStreamzName[] =
    "/tensorflow/core/tf2xla/api/v2/phase2_compilation_status";

using ::tensorflow::monitoring::testing::CellReader;

static constexpr char kFailedLegalizationStreamz[] =
    "/tensorflow/core/tf2xla/mlir_second_phase_failed_legalization_op_count";
static constexpr char kNonStaticOpStreamz[] =
    "/tensorflow/core/tf2xla/mlir_second_phase_non_static_op_count";
static constexpr char kNonStaticOpSkipStreamz[] =
    "/tensorflow/core/tf2xla/mlir_second_phase_non_static_op_skip_count";

// the SessionInitializerOp.
inline constexpr StringRef kTfSavedModelInitializerTypeAttr =
    "tf_saved_model.initializer_type";

// Indicates that the initializer corresponds to the restore op.
inline constexpr StringRef kTfSavedModelInitializerRestoreType = "restore_op";

// Indicates that the initializer corresponds to the init op.
inline constexpr StringRef kTfSavedModelInitializerInitType = "init_op";