}

// CHECK: func private @[[LAUNCH]]
// CHECK-SAME: (%[[LAUNCH_ARG_0:[a-z0-9]*]]: tensor<?xi32>) -> tensor<?xi32>
// CHECK: %[[B_OUTPUT:[0-9]*]] = "tf.B"(%[[LAUNCH_ARG_0]])
// CHECK: return %[[B_OUTPUT]]

// -----

// Tests that multiple `tf_device.launch` that depend on each other are
// correctly handled.

// CHECK-LABEL: func @multiple_launches
// CHECK-SAME: (%[[ARG_0:[a-z0-9]*]]: tensor<?xi32>)

  // CHECK: "tf.Iterator"
  // CHECK-SAME: output_shapes{{.*}}200x28x28x1{{.*}}200x10
  // CHECK-SAME: output_types = [f32, f32]
  // CHECK-SAME: shared_name = "_iterator1"
  %0 = "tf.AnonymousIteratorV3"() {output_shapes = [
    #tf_type.shape<200x28x28x1>,
    #tf_type.shape<200x10>], output_types = [f32, f32]} : () -> tensor<!tf_type.resource>
  // CHECK: "tf.Iterator"
  // CHECK-SAME: shared_name = "_iterator2"

}

// CHECK-LABEL: func private @then_removed_const_args_0
// CHECK-SAME: ([[x:%.*]]: tensor<i32>)
// CHECK: [[const:%.*]] = "tf.Const"
// CHECK-SAME: value = dense<10> : tensor<i32>
// CHECK: [[r:%.*]] = "tf.StatefulPartitionedCall"([[x]], [[const]])
// CHECK-SAME: f = @then}
// CHECK: return [[r]]

// CHECK-LABEL: func private @else_removed_const_args_0
// CHECK-SAME: ([[x:%.*]]: tensor<i32>)
// CHECK: [[const:%.*]] = "tf.Const"

  }
}

// CHECK-LABEL: func @conv2d_unmatching_unit
// CHECK: "tf.PartitionedCall"
// CHECK-SAME: f = @composite_conv2d_fn_1
// Check that the `_tfl_quant_trait` attribute exists since the unit is not in `unit_wise_quantization_specs`.
// CHECK-SAME: _tfl_quant_trait = "fully_quantizable"
// CHECK-SAME: loc(callsite("Model/conv2d@conv2d_unmatching_unit"("Conv2D") at "QuantizationUnit({{.*}})"))

// -----

// disabled by having `_quantization_method = "no_quantization { }"`.
// DISABLE-ALL-DOT-GENERAL-SAME: _entry_function = @composite_dot_general_fn_1
// DISABLE-ALL-DOT-GENERAL-SAME: _original_entry_function
// DISABLE-ALL-DOT-GENERAL-SAME: _quantization_method = "no_quantization { }"
// DISABLE-ALL-DOT-GENERAL-SAME: _tfl_quant_trait = "fully_quantizable"

// DISABLE-ALL-DOT-GENERAL: return %[[XLA_CALL_MODULE:.+]] : tensor<1x1x64xf32>

# CHECK:      func {{@[a-zA-Z_0-9]+}}(
# CHECK-SAME:   [[ARG0:%.*]]: tensor<3x1xf32> {tf_saved_model.index_path = ["x"]},
# CHECK-SAME:   [[ARG1:%.*]]: tensor<!tf_type.resource<tensor<1x3xf32>>> {tf_saved_model.bound_input = @[[VAR]]})
# CHECK-SAME:             -> (tensor<3x3xf32> {tf_saved_model.index_path = ["r"]})
# CHECK-SAME: attributes {{.*}} tf_saved_model.exported_names = ["key"]

// CHECK-LABEL: func @dot_int4
// CHECK-SAME: (%[[ARG_0:.*]]: tensor<?x3xf32>) -> tensor<?x2xf32>
func.func @dot_int4(%arg0: tensor<?x3xf32>) -> tensor<?x2xf32> {
  // CHECK: %[[cst:.*]] = stablehlo.constant
  // CHECK: %[[q1:.*]] = "quantfork.qcast"(%[[cst]])
  // CHECK-SAME: quant.uniform<i8:f32, 0.0040316890267764818:127>
  // CHECK: %[[dq1:.*]] = "quantfork.dcast"(%[[q1]])
  // CHECK-SAME: quant.uniform<i8:f32, 0.0040316890267764818:127>

// CHECK-LABEL: func @testSingleConj
// CHECK-SAME:  ([[ARG0:%.+]]: tensor<complex<f32>>)
func.func @testSingleConj(%arg0: tensor<complex<f32>>) -> tensor<complex<f32>> {
  // CHECK: [[CONJ:%.+]] = "tf.Conj"([[ARG0]])
  %0 = "tf.Conj"(%arg0) : (tensor<complex<f32>>) -> tensor<complex<f32>>
  // CHECK: return [[CONJ]]
  func.return %0: tensor<complex<f32>>
}

// CHECK-LABEL: func @testDoubleConj
// CHECK-SAME:  ([[ARG0:%.+]]: tensor<complex<f32>>)

# of the feed.
#
# CHECK-LABEL: func @main
# CHECK-SAME:  (%[[ARG_0:.*]]: tensor<f32>, %[[ARG_1:.*]]: tensor<f32>) -> (tensor<f32>, tensor<f32>, tensor<*xf32>)
# CHECK-SAME:  control_outputs = ""
# CHECK-SAME:  inputs = "z:1,z:2"
# CHECK-SAME:  outputs = "z:2,z:1,a:0"
# CHECK:           %{{.*}}, %[[ASSIGN_ADD_CTRL:.*]] = tf_executor.island wraps "tf.AssignAdd"

  // FB-SAME:  -8.0, -9.0, -10.0, -11.0, -12.0, -13.0, -14.0, -15.0, -16.0,
  // FB-SAME:  -17.0, -18.0, -19.0, -20.0, -21.0, -22.0, -23.0, -24.0, -25.0,
  // FB-SAME:  -26.0, -27.0, -28.0, -29.0, -30.0, -31.0, -32.0, -33.0, -34.0,
  // FB-SAME:  -35.0, -36.0, -37.0, -38.0, -39.0 ],
  // FB-NEXT:        max: [ 0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0,