// RUN: tac-opt-all-backends -tfl-device-transform-nnapi %s -split-input-file -verify-diagnostics | FileCheck %s

func.func @mean_4d_keepdim(%arg0: tensor<1x48x48x512xf32>) -> tensor<1x1x1x512xf32> {
  %cst = arith.constant dense<[1, 2]> : tensor<2xi32>
  %0 = "tfl.mean"(%arg0, %cst) {keep_dims = true} : (tensor<1x48x48x512xf32>, tensor<2xi32>) -> tensor<1x1x1x512xf32>
  func.return %0 : tensor<1x1x1x512xf32>
}

// Creates a pass that extracts outside compilation (Host ops inside device
// cluster) at head/tail of Device cluster to run before/after XLA computation.
std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
CreateExtractHeadTailOutsideCompilationPass();

// Creates a pass that extract outside compilation (Host ops inside cevice
// cluster) ops to a separate parallel_execute region to run on CPU.

    return std::nullopt;

  if (!HasValidHardwareTarget(op)) return std::nullopt;

  auto device = op->getAttrOfType<mlir::StringAttr>(mlir::TFL::tac::kDevice);
  if (device == nullptr) return std::nullopt;

  llvm::StringRef device_name_str = device.getValue();
  return device_name_str.str();
}

std::optional<std::vector<float>> GetPerDeviceCosts(

                                       mlir::Builder* builder) {
  // Parse GPU device compute capability from physical device description.
  static auto* r = new llvm::Regex("compute capability: ([0-9]+)\\.([0-9]+)");

  llvm::SmallVector<llvm::StringRef, 3> cc;
  if (r->match(device.attributes().physical_device_desc(), &cc)) {
    return mlir::TF::GpuDeviceMetadata::get(

    %outputs, %control = tf_executor.island wraps "tf.GuaranteeConst"(%arg1) {T = f32, device = ""} : (tensor<f32>) -> tensor<f32>
    %outputs_0, %control_1 = tf_executor.island wraps "tf.GuaranteeConst"(%arg2) {T = f32, device = ""} : (tensor<f32>) -> tensor<f32>
    %control_2 = tf_executor.island wraps "tf.NoOp"() {_pivot_for_cluster = "cluster", device = ""} : () -> ()

      %outputs_11, %control_12 = tf_executor.island wraps "tf.Placeholder"() {device = "", dtype = "tfdtype$DT_FLOAT", name = "y", shape = "tfshape$dim { }"} : () -> tensor<0xf32>

    %1 = "tf.Relu"(%0) {device = "tpu0"} : (tensor<?xi32>) -> tensor<?xi32>
    %2 = "tf.Add"(%0, %1) {device = "tpu0"} : (tensor<?xi32>, tensor<?xi32>) -> tensor<?xi32>
    %3 = "tf.Relu"(%arg1) {device = "tpu1"} : (tensor<?xi32>) -> tensor<?xi32>
    %4 = "tf.Add"(%3, %arg1) {device = "tpu1"} : (tensor<?xi32>, tensor<?xi32>) -> tensor<?xi32>
    %5 = "tf.Relu"(%4) {device = "tpu0"} : (tensor<?xi32>) -> tensor<?xi32>

  // CHECK: {{%.*}} = tfrt_fallback_async.executeop {{.*}} device("/device:CPU:0") "tf.MatMul"
  // CHECK-SAME: {T = f32, transpose_a = false, transpose_b = false}

    %2:3 = "tf.UnpackHyp"(%arg1) {_tpu_replicate = "cluster_Fn", _xla_outside_compilation = "0", device = "", max_seq_length = 16 : i64} : (tensor<*x!tf_type.string>) -> (tensor<*xi32>, tensor<*xi32>, tensor<*xf32>)
    %3 = "tf.Reshape"(%2#2, %1) {_tpu_replicate = "cluster_Fn", _xla_outside_compilation = "0", device = ""} : (tensor<*xf32>, tensor<*xi32>) -> tensor<*xf32>

  %dense_default_1 = "tf.Const"() {device = "/device:CPU:0", dtype = f32, value = dense<[]> : tensor<0xf32>} : () -> tensor<0xf32>
  %dense_keys = "tf.Const"() {device = "/device:CPU:0", dtype = !tf_type.string, value = dense<""> : tensor<2x!tf_type.string>} : () -> tensor<2x!tf_type.string>