func.return %0 : tensor<?xi32>
}

// -----

func.func @constant_tensor_invalid_2(%arg0: vector<1xi32>) -> tensor<1xf32> {
    // expected-error@+1 {{input and output should have same shape and element type}}
  %0 = "tfr.constant_tensor"(%arg0) : (vector<1xi32>) -> tensor<1xf32>
  func.return %0 : tensor<1xf32>
}

// -----

func.func @constant_tensor_invalid_3(%arg0: vector<1xi32>) -> tensor<1x1xi32> {

func.func @add_with_activation_transpose_broadcasted_rhs(%arg0: tensor<1x3x3x4xf32>) -> tensor<1x4x3x3xf32> {
  %0 = stablehlo.constant dense<2.000000e+00> : tensor<4xf32>
  %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<4xf32>) -> tensor<1x4x3x3xf32>
  %2 = stablehlo.transpose %arg0, dims = [0, 3, 1, 2] : (tensor<1x3x3x4xf32>) -> tensor<1x4x3x3xf32>
  %3 = stablehlo.add %2, %1 : tensor<1x4x3x3xf32>

    %2 = "tf.MatMul"(%arg2, %arg1) {} : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
    %3 = "tf.AddV2" (%arg1, %arg1)  : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
    %4 = "tf.Identity"(%1) {device = ""} : (tensor<i32>) -> tensor<i32>
    %5 = "tf.Identity"(%3) {device = ""} : (tensor<*xf32>) -> tensor<*xf32>
    %6 = "tf.Identity"(%2) {device = ""} : (tensor<*xf32>) -> tensor<*xf32>

}

// -----

// CHECK-LABEL: func @convolution_add_add_f32
func.func @convolution_add_add_f32(
    %lhs: tensor<?x3x2x1xf32>, %rhs: tensor<2x1x1x1xf32>,
    %zp_offset: tensor<?x2x2x1xf32>, %bias: tensor<1xf32>
  ) -> tensor<?x2x2x1xf32> {
  // CHECK-DAG: %[[conv:.*]] = mhlo.convolution
  // CHECK-DAG: %[[combined:.*]] = chlo.broadcast_add %[[conv:.*]], %[[zp_offset:.*]]

// CHECK-LABEL: QuantizeCustomOp
// CustomOp-LABEL: QuantizeCustomOp
func.func @QuantizeCustomOp(%arg0: tensor<1x1x1x1xf32>) -> (tensor<*xf32>, tensor<*xf32>, tensor<*xf32>) attributes {tf.entry_function = {inputs = "input", outputs = "custom_op"}} {
  %0 = "quantfork.stats"(%arg0) {layerStats = dense<[0.000000e+00, 2.550000e+02]> : tensor<2xf32>} : (tensor<1x1x1x1xf32>) -> tensor<1x1x1x1xf32>
  %w_1 = arith.constant dense<127.0> : tensor<4096x1x1x1xf32>

func.func @no_qdq_for_non_weight_constant(%arg0: tensor<1x2xf32>, %arg1: tensor<2x3xf32>) -> tensor<1x3xf32> attributes {tf._original_func_name = "main_0"} {
  %cst = "tf.Const"() {value = dense<4.000000e-02> : tensor<3xf32>} : () -> tensor<3xf32>
  %0 = "tf.XlaCallModule"(%arg0, %arg1, %cst) {
    Sout = [#tf_type.shape<1x3>], _entry_function = @composite_dot_general_with_bias_fn,
    _original_entry_function = "composite_dot_general_with_bias_fn",

  // CHECK-NEXT: %[[UPDATE_SLICE:.*]] = "tf.Reshape"(%[[PUSHVAL]], %[[UPDATE_SHAPE]]) : (tensor<f32>, tensor<1xi32>) -> tensor<1xf32>
  // CHECK-NEXT: %[[UPDATE:.*]] = "tf.XlaDynamicUpdateSlice"(%[[READ_VAL]], %[[UPDATE_SLICE]], %[[READ_SIZE]]) : (tensor<10xf32>, tensor<1xf32>, tensor<1xi32>) -> tensor<10xf32>
  // CHECK-NEXT: "tf.AssignVariableOp"(%[[BUFFER]], %[[UPDATE]]) : (tensor<!tf_type.resource<tensor<10xf32>>>, tensor<10xf32>) -> ()

    tensor type, the shape shouldn't be changed during the conversion.

    Example:

    ```mlir
    %1 = tfr.constant_tensor(%0) : f32 -> tensor<f32>
    %3 = tfr.constant_tensor(%2) : vector<1xf32> -> tensor<1xf32>
    ```
  }];

  let arguments = (ins TFR_AllAttrTypes:$arg);

  let results = (outs TFR_singleTensorType:$out);

  let hasCanonicalizer = 1;

  let hasVerifier = 1;
}

// CustomOpNotWeightOnly-LABEL: QuantizeCustomOp
func.func @QuantizeCustomOp(%arg0: tensor<1x1x1x1xf32>) -> tensor<*xf32> attributes {tf.entry_function = {inputs = "input", outputs = "custom_op"}} {
  %0 = "quantfork.stats"(%arg0) {layerStats = dense<[0.000000e+00, 2.550000e+02]> : tensor<2xf32>} : (tensor<1x1x1x1xf32>) -> tensor<1x1x1x1xf32>
  %w = arith.constant dense<127.0> : tensor<1024x1x1x1xf32>

// (i.e., after converting reference types to their corresponding TensorFlow or
// standard types). Also, this allows compatible types so it is legal to have
// tensor<*xf32> and tensor<4xf32> types.
def TF_SameOperandsAndResultTypeResolveRef : TraitList<
  InferTensorType.traits #
  [
    NativeOpTrait<"TF::SameOperandsAndResultTypeResolveRef">
  ]>;