// when loading and unloading large numbers of SavedModelBundles.
// TODO(sukritiramesh): Increase run iterations and move outside of the test
// suite.
TEST_F(LoaderTest, ResourceLeakTest) {
  SavedModelBundle bundle;
  SessionOptions session_options;
  RunOptions run_options;

  const string export_dir =
      io::JoinPath(testing::TensorFlowSrcRoot(), kTestDataSharded);

  // Run the graph
  Run({0, 0});
  ExpectOutputValue(0, 8);
  ExpectOutputValue(1, 3);
}

TEST_F(CApiWhileLoopTest, UnsetCondOutput) {
  Init(1);
  params_->body_outputs[0] = params_->body_inputs[0];
  ExpectError(TF_INVALID_ARGUMENT,
              "TF_WhileParams `cond_output` field isn't set");
}

TEST_F(CApiWhileLoopTest, WrongCondOutputType) {
  Init(1);
  params_->cond_output = params_->cond_inputs[0];

      %0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0], precision = [DEFAULT, DEFAULT] : (tensor<1x1024xf32>, tensor<1024x3xf32>) -> tensor<1x3xf32>
      return %0 : tensor<1x3xf32>
    }
  }
)mlir";

TEST_F(LiftAsFunctionCallTest, LiftedFunctionSucceeds) {
  OwningOpRef<ModuleOp> module_op = ParseModuleOpString(kModuleLifted);
  ASSERT_TRUE(module_op);

  auto composite_dot_general_fn =

  key.set_device_type(device_type.type_string());
  key.set_prefix(persistence_prefix);
  key.set_compiled_using_pjrt(compiled_using_pjrt);
  return key;
}

TEST_F(DeviceExecutionPersistorTest, PersistCacheDirNotSet) {
  XlaDeviceExecutablePersistor::Config config(
      /*persistent_cache_directory=*/"",
      /*disable_strict_signature_checks=*/false,

  OwningOpRef<ModuleOp> module_;
  llvm::SourceMgr source_manager_;
};

TEST_F(Tf2XlaRewriterTest, LegalizesOpWithTf2xlaHloImporter) {
  TF_EXPECT_OK(LegalizeModule());

  int num_tuple_ops = 0;
  module_->walk([&num_tuple_ops](TupleOp tuple_op) { num_tuple_ops += 1; });

  EXPECT_EQ(num_tuple_ops, 0);
}

TEST_F(Tf2XlaRewriterTest, ImportsXlaComputationIntoModule) {
  TF_ASSERT_OK(CreateMlirModule());

    using Rematerializer::AddTensor;
    using Rematerializer::AddUse;
    using Rematerializer::DelUse;
    using Rematerializer::Remat;
  };
  TestableRematerializer r_;
};

TEST_F(RematTest, TensorUseSimple) {
  for (int i = 0; i < 6; ++i) {
    r_.AddOperation(/*is_stateful=*/false);
    r_.AddTensor(/*size=*/1 << i);
  }

  r_.AddUse(/*ioperation=*/2, /*itensor=*/2);

}

TEST_F(FreezeTest, GraphDefWithoutDependentVariables) {
  TestFreezeGraphWithoutDependentVariables(false);
}

TEST_F(FreezeTest, GraphDefWithoutDependentResourceVariables) {
  TestFreezeGraphWithoutDependentVariables(true);
}

TEST_F(FreezeTest, GraphDefWithDependentVariables) {
  TestFreezeGraphWithDependentVariables(false);
}

class XlaCompilerOptionsTest : public ::testing::Test {
 protected:
  void SetUp() override {
    tensorflow::GetXlaDeviceFlags()->tf_xla_enable_xla_devices = true;
  }

  DeviceSetup device_setup_;
};

TEST_F(XlaCompilerOptionsTest, PjRtOptionsXlaDevice) {
  device_setup_.AddDevicesAndSetUp({DEVICE_XLA_GPU});
  Device* device = device_setup_.GetDevice(DEVICE_XLA_GPU);

    auto module_op_ref =
        mlir::parseSourceString<mlir::ModuleOp>(module_op_str, &ctx_);
    EXPECT_TRUE(module_op_ref);
    return module_op_ref;
  }

  Env* env_{};
  mlir::MLIRContext ctx_{};
};

TEST_F(SaveVariablesToCheckpointTest, VariableSavedToCheckpoint) {
  constexpr absl::string_view kModuleCode = R"mlir(
    module attributes {tf_saved_model.semantics} {

using ::testing::ElementsAre;
using ::testing::HasSubstr;
using ::testing::Key;
using ::testing::SizeIs;
using ::tsl::testing::StatusIs;

class CalibrationStatisticsSaverTest : public OpsTestBase {};

TEST_F(CalibrationStatisticsSaverTest, MissingOutputPath) {
  std::vector<std::string> ids{"1"};
  std::vector<int32_t> calibration_methods{
      CalibrationOptions::CALIBRATION_METHOD_AVERAGE_MIN_MAX};