/*shape_determination_fns=*/{}, arg_shapes,
                     &arg_core_mapping, &per_core_arg_shapes, client)
                     .status());
  }
};

TEST_F(LegalizeTFQuantTest, LegalizesModuleWithTFUniformQuantization) {
  constexpr char mlir_module_string[] = R"mlir(
  module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {

  spec->biases_params[2] = {{0, 1}, quant::GetUniformQuantizedTypeForBias};
  for (const auto& [key, value] : spec->coeff_op_quant_dim) {
    spec->quantizable_operands.insert(key);
  }
  return spec;
}

TEST_F(ApplyQuantizationParamsPropagationTest,
       ConstsUsedMultipleTimesAreDuplicated) {
  const OwningOpRef<ModuleOp> module_op_ref =
      ParseModuleOpString(kModuleTFLite);

    return run_result;
  }

 private:
  mlir::MLIRContext context_;
  mlir::OwningOpRef<ModuleOp> module_;
  std::unique_ptr<mlir::PassManager> pm_;
};

TEST_F(InputLoweringMetricsPassTest, CountsNoDynamicOps) {
  static constexpr char kMlirModuleStr[] = R"(
  module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {

    }
    return absl::OkStatus();
  }

  DialectRegistry registry_;
  MLIRContext context_;
  OwningOpRef<mlir::ModuleOp> mlir_module_;
};

TEST_F(TpuMetadataUtilsTest, SingleDevice) {
  TF_ASSERT_OK_AND_ASSIGN(auto cluster_func_ops,
                          GetClusterFuncOps("basic_cluster.mlir"));

    for (const std::string& file : files) {
      TF_ASSERT_OK(env_->DeleteFile(file));
    }
  }

  tsl::Env* env_;
  std::string test_dir_;
  std::unique_ptr<mlir::MLIRContext> ctx_;
};

TEST_F(EnableIrPrintingTest, PassSuccessfullyRuns) {
  mlir::PassManager pm = {ctx_.get()};
  pm.addPass(CreateNoOpPass());
  pm.addNestedPass<mlir::func::FuncOp>(mlir::createCanonicalizerPass());

  func::FuncOp fused_lstm_func_;
  func::FuncOp fused_lstm_func_cifg_;
  func::FuncOp fused_ln_lstm_func_;
  std::unique_ptr<mlir::MLIRContext> context_;
  std::unique_ptr<mlir::Builder> builder_;
};

TEST_F(LstmUtilsTest, ConvertLSTMCellSimple) {
  mlir::TFL::ConvertLSTMCellSimpleToFusedLSTM convert(fused_lstm_func_);

  auto result = convert.RewriteFunc();
  EXPECT_FALSE(failed(result));
  fused_lstm_func_.dump();

    readonly_model_ = input_model_->GetModel();
    model_ = UnPackFlatBufferModel(*readonly_model_);
  }
};

// TODO(b/327796566): re-enable after the bug is fixed
// TEST_F(QuantizeConvNoBiasModelTest, QuantizationSucceeds) {
//   auto status = QuantizeModelAllOperators(
//       &model_, TensorType_INT8, TensorType_INT8,
//       /*allow_float=*/false, TensorType_INT8, &error_reporter_,

  tensorflow::core::RefCountPtr<DeviceContext> device_context_;
  tensorflow::Allocator* host_allocator_;
  tensorflow::Allocator* device_allocator_;
};

#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
TEST_F(DeviceContextTest, TestXlaGpuRoundTripTransferWithDeviceApi) {
  SetDevice(DEVICE_XLA_GPU);
  tensorflow::Tensor origin_cpu_tensor(host_allocator_, tensorflow::DT_FLOAT,

// class.  For example:
//
// class MyClass {
//  private:
//   void MyMethod();
//   FRIEND_TEST(MyClassTest, MyMethod);
// };
//
// class MyClassTest : public testing::Test {
//   // ...
// };
//
// TEST_F(MyClassTest, MyMethod) {
//   // Can call MyClass::MyMethod() here.
// }

#define FRIEND_TEST(test_case_name, test_name)\
friend class test_case_name##_##test_name##_Test

        flr, fld, shape_inference_graphs, has_outside_compilation);
  }

 private:
  std::unique_ptr<DeviceMgr> device_mgr_;
  std::unique_ptr<ProcessFunctionLibraryRuntime> pflr_;
};

TEST_F(ExtractOutsideCompilationForFunctionTest, Basic) {
  // Build the XLA computation func.
  // "const0"
  // "identity0" = "const0" (outside compilation cluster "0")