// RUN: tf-opt "--tfxla-device-specific-transforms=device-type=XLA_GPU_JIT" -verify-diagnostics -split-input-file %s | FileCheck -dump-input=fail %s

module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 1399 : i32}} {

// CHECK-LABEL: stateless_op
func.func @stateless_op() -> tensor<i32> {
  // CHECK: %cst = "tf.Const"() <{value = dense<1> : tensor<i32>}> : () -> tensor<i32>

// RUN: tf-opt "--tfxla-device-specific-transforms=device-type=XLA_CPU_JIT" -verify-diagnostics -split-input-file %s | FileCheck -dump-input=fail %s

module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 1399 : i32}} {

// CHECK-LABEL: stateless_op
func.func @stateless_op() -> tensor<i32> {
  // CHECK: %cst = "tf.Const"() <{value = dense<1> : tensor<i32>}> : () -> tensor<i32>

  pm.addPass(mlir::createCanonicalizerPass());
  pm.addPass(mlir::createSymbolDCEPass());

  pm.addPass(mlir::TF::CreateTFShapeInferencePass());

  llvm::StringRef tf2xla_fallback_device_type = "XLA_CPU_JIT";
  pm.addPass(mlir::mhlo::createLegalizeTFPass(
      /*legalize_chlo=*/true, tf2xla_fallback_device_type,
      /*prefer_tf2xla=*/false));

  pm.addPass(mlir::TF::CreateStripTfAttributesPass());

// RUN: tf-opt "-xla-legalize-tf=use-tf2xla-fallback=true device-type=XLA_CPU_JIT" -verify-diagnostics %s | FileCheck --check-prefix SUPPORTED_FALLBACK_DEVICE %s
// RUN: tf-opt "-xla-legalize-tf=use-tf2xla-fallback=true" %s | FileCheck --check-prefix UNSPECIFIED_FALLBACK_DEVICE %s
// RUN: tf-opt "-xla-legalize-tf=use-tf2xla-fallback=true device-type=INVALID_DEVICE_TYPE" %s | FileCheck --check-prefix UNSUPPORTED_FALLBACK_DEVICE %s

// We run this test four times:

  explicit Tf2XlaRewriterTestPeer(mlir::Operation* op)
      : op_builder_(op),
        empty_rewriter_(op_builder_),
        tf2xla_rewriter_(op, empty_rewriter_,
                         /*device_type=*/"XLA_CPU_JIT") {}

  absl::StatusOr<TupleOp> ImportXlaComputationIntoModule(
      XlaComputation& computation) {
    return tf2xla_rewriter_.ImportXlaComputation(computation);
  }

 private:
  OpBuilder op_builder_;

    std::string device_name;

    // The index of the device.
    int device_ordinal = -1;

    // The name of the compilation device, also referred to as jit_device_type.
    // (e.g., "XLA_CPU_JIT");
    std::string compilation_device_name;

    // A vector of ShapeDeterminationFn (i.e., a bundle of LayoutSelectionFn,
    // ShapeRepresentationFn). Each bundle describes how the on-host shapes of

  mhlo::PopulateLegalizeTfPatterns(context, &patterns);
  TF::PopulateTFLoweringBeforeHLOPatterns(context, &patterns);
  mhlo::Tf2XlaTypeConverter converter;
  mhlo::PopulateLegalizeTfWithTf2XlaPatterns(
      "XLA_CPU_JIT", patterns, context, converter, /*prefer_tf2xla=*/false);
  stablehlo::StablehloToHloTypeConverter hlo_converter;
  chlo::populateChloToHloPatterns(context, &hlo_converter, &patterns);

static constexpr char kMlirToGraphCompilationCheckName[] =
    "mlir-to-graph-compilation-check";
// Use CPU arbitrarily in order to check that a graph compiles at all
static constexpr char kArbitraryDeviceName[] = "XLA_CPU_JIT";

namespace {
inline absl::string_view StringRefToView(llvm::StringRef ref) {
  return {ref.data(), ref.size()};
}
}  // namespace

static OwningOpRef<mlir::ModuleOp> GraphdefToMlirTranslateFunction(

      ParseArgumentShapes(mlir::StringRefToView(input_shapes), arg_shapes);
  if (!args_status.ok()) {
    LOG(ERROR) << args_status;
    return mlir::failure();
  }

  auto device_type = "XLA_CPU_JIT";
  llvm::MutableArrayRef<std::unique_ptr<mlir::Pass>>
      custom_legalization_passes{};
  XlaCompilationResult compilation_result;
  auto compilation_status =
      via_builder

// RUN: tf-opt "-xla-legalize-tf=device-type=XLA_CPU_JIT legalize-chlo=false use-tf2xla-fallback=true prefer-tf2xla=true" %s | FileCheck %s
// RUN: tf-opt "-xla-legalize-tf=device-type=XLA_CPU_JIT legalize-chlo=false prefer-tf2xla=true" %s | FileCheck --check-prefix NOFALLBACK %s

module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {

// CHECK-LABEL: @abs
func.func @abs(%arg0: tensor<2xf32>) -> tensor<2xf32> {