// CHECK-NEXT:     signature_defs: [ ]
// CHECK-NEXT:   }

// IMPORT: "tfl.fake_quant"(%arg0) <{max = 1.400000e+00 : f32, min = 3.000000e-01 : f32, narrow_range = false, num_bits = 6 : i32}>

  %0 = "tfl.fake_quant"(%arg0) {num_bits = 6 : i32, narrow_range = false, min = 0.3:f32, max = 1.4:f32} : (tensor<4 x f32>) -> tensor<4 x f32>
  func.return %0 : tensor<4xf32>

  %cst = "tf.Const"() <{value = dense<1> : tensor<i32>}> {device = ""} : () -> tensor<i32>
  %0 = "tf.MultiDeviceIteratorFromStringHandle"(%arg0) <{output_shapes = [#tf_type.shape<>], output_types = [!tf_type.string]}> {device = ""} : (tensor<!tf_type.string>) -> tensor<!tf_type.resource>
  %1 = "tf.MultiDeviceIteratorGetNextFromShard"(%0, %cst, %arg1) {device = ""} : (tensor<!tf_type.resource>, tensor<i32>, tensor<i64>) -> tensor<!tf_type.string>

func.func @shape() {
  %size = "tf.Const"() {value = dense<5> : tensor<i32>} : () -> tensor<i32>
  // CHECK: tf.TensorArrayV3
  // CHECK-SAME: element_shape = #corert.shape<*>
  %ta:2 = "tf.TensorArrayV3"(%size) {device = "/device:CPU:0", dtype = f32, element_shape = #tf_type.shape<*>, dynamic_size = false, clear_after_read = true, identical_element_shapes = true, tensor_array_name = "ta"} : (tensor<i32>) -> (tensor<!tf_type.resource>, tensor<f32>)
  func.return
}

        self.assertIn('stablehlo.add %arg0, %cst : tensor<4xf32>', str(module))

  def test_tf_mlir_to_stablehlo(self):
    assembly = """
      module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
        func.func @main(%arg0 : tensor<?xf32>) -> tensor<?xf32> {
          %cst = "tf.Const"() {value = dense<1.0> : tensor<f32>} : () -> tensor<f32>

        /* prefer_tf2xla=*/false));
  };
}

TEST(XlaLegalizeTest, IllegalOp) {
  constexpr char kMlirIllegalOpStr[] = R"(
  module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
    func.func @main() -> tensor<1xi32> {
      %0 = "tf.DoesntExist"() : () -> tensor<1xi32>
      func.return %0 : tensor<1xi32>
    }
  })";
  CellReader<int64_t> legalize_failure_count(

    fn.walk([&](Operation* op) {
      if (!llvm::isa<TFL::ConstOp, TFL::QConstOp, arith::ConstantOp>(op))
        return;

      // We only fold int32/int64 for Const and i32 for QConst if not specify
      // all constants flag. (Since they're more like "configs" or i32 biases.)
      // We will fold every const ops (and q_const ops) if we speicfy the
      // fold_all_constants_flag.
      if (!fold_all_constants_flag_) {

  }
  if (AttrType content = mlir::dyn_cast<AttrType>(attr)) {
    *out_attr = content;
    return true;
  } else {
    return false;
  }
}

// Changes a DenseIntElementsAttr **containing I64** elements to an I32 Vector.
bool DenseI64AttrToI32Vector(const DenseIntElementsAttr& dense_attr,
                             std::vector<int32_t>* out_vec);

// Gets boolean from composite attrs if it exists.

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

static constexpr char kMlirModuleStr[] = R"(
  module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
    func.func @main() -> tensor<1xi32> {
      %0 = "tf.Const"() {value = dense<1000> : tensor<1xi32>} : () -> tensor<1xi32>
      func.return %0 : tensor<1xi32>
    }
  })";

  func.return %0 : tensor<40x40xf32>

// CHECK-LABEL: func @main(%arg0: tensor<40x37xf32>, %arg1: tensor<40x37xf32>) -> tensor<40x40xf32>
// CHECK:      %[[CONST:[0-9]+]] = "tfl.external_const"() <{buffer_index = 3 : i32}> : () -> tensor<40xf32>
// CHECK-NEXT: %[[FULL:[0-9]+]]:2 = "tfl.fully_connected"(%arg0, %arg1, %[[CONST]]) <{asymmetric_quantize_inputs = false, fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"}>

  %3 = "tfl.pack"(%1, %2) {axis = 0 : i32, per_device_costs = {CPU = 2.0 : f32, GPU = -1.0 : f32}, values_count = 2 : i32, tac.device = "CPU"} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2x1xf32>
  func.return %3 : tensor<2x1xf32>
})";
  const std::string kExpectedFB = CreateRuntimeMetadata();