// CHECK:  [[VAL_1:%.*]] = "tfl.reshape"(%2, %[[CST]]) {tac.device = "GPU",  tac.inference_type = "FLOAT"} : (tensor<1xf32>, tensor<4xi32>) -> tensor<1x1x1x1xf32>
// CHECK:  [[VAL_2:%.*]] = "tfl.concatenation"([[VAL_0]], [[VAL_1]]) <{axis = 3 : i32, fused_activation_function = "NONE"}> {tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<1x1x1x1xf32>, tensor<1x1x1x1xf32>) -> tensor<1x1x1x2xf32>

    return %5 : tensor<1x3x1x1xf32>
  }
  func.func private @composite_gather_fn_1(%arg0: tensor<1x3x1x1xf32>, %arg1: tensor<1xi32>, %arg2: tensor<i32>) -> tensor<1x3x1x1xf32> attributes {tf_quant.composite_function} {

func.func @depth_to_space(%arg0: tensor<1x1x1x4xf32>) -> tensor<1x2x2x1xf32> {
  %0 = "tf.DepthToSpace"(%arg0) {block_size = 2: i64,  data_format = "NHWC"}: (tensor<1x1x1x4xf32>) -> tensor<1x2x2x1xf32>
  func.return %0 : tensor<1x2x2x1xf32>
// CHECK: %[[CUSTOM_0:.*]] = "tfl.custom"(%arg0) <{custom_code = "FlexDepthToSpace", custom_option = #tfl<const_bytes : "{{.*}}">}> : (tensor<1x1x1x4xf32>) -> tensor<1x2x2x1xf32>
// CHECK: return %[[CUSTOM_0]] : tensor<1x2x2x1xf32>
}

    // CHECK-NEXT: %[[CMP:.*]] = mhlo.compare GT, %[[INP]], %[[ZERO]], NOTYPE : (tensor<1x4x4x3xf32>, tensor<1x4x4x3xf32>) -> tensor<1x4x4x3xi1>
    // CHECK-NEXT: %[[RES:.*]] = mhlo.select %[[CMP]], %[[INP]], %[[LEAKY]] : tensor<1x4x4x3xi1>, tensor<1x4x4x3xf32>
    // CHECK-NEXT: return %[[RES]] : tensor<1x4x4x3xf32>
    %0 = "tf.LeakyRelu"(%arg0) {alpha = 2.000000e-01 : f32, device = ""} : (tensor<1x4x4x3xf32>) -> tensor<1x4x4x3xf32>

    %dimension = "tf.Const"() { value = dense<3> : tensor<1xi64> } : () -> tensor<1xi64>
    %6 = "tf.Sum"(%3, %dimension) { keep_dims = true }: (tensor<1x3x2x2xf32>, tensor<1xi64>) -> tensor<1x3x2x1xf32>
    return %5, %6 : tensor<1x3x2x2xf32>, tensor<1x3x2x1xf32>
  }

// CHECK-LABEL: func @conv_with_multiple_uses
// CHECK: %[[div:.*]] = "tf.Div"(%arg0
// CHECK: %[[add:.*]] = "tf.AddV2"(%[[div]]
// CHECK: %[[maximum:.*]] = "tf.Maximum"(%[[add]]

    %1 = stablehlo.constant dense<1.000000e+03> : tensor<1x1x1x1xf32>  // Input inverse scale.
    %2 = stablehlo.constant dense<-128> : tensor<1x1x1x1xi8>  // Input zero point.
    %3 = stablehlo.constant dense<1> : tensor<3x3x4x4xi8>  // Quantized filter tensor.
    %4 = stablehlo.constant dense<3.000000e+03> : tensor<1x1x1x4xf32>
    %5 = stablehlo.constant dense<4.000000e+03> : tensor<1x1x1x1xf32>  // Output inverse scale.

// RUN: tf-opt -split-input-file -verify-diagnostics -tf-einsum %s | FileCheck %s

func.func @unary_einsum_reduce_sum_transpose(%arg0: tensor<3x4x5x6xf32>) -> tensor<3x5x4xf32> {
  %0 = "tf.Einsum"(%arg0) {T = "tfdtype$DT_FLOAT", equation = "...gse->...sg"}: (tensor<3x4x5x6xf32>) -> tensor<3x5x4xf32>
  func.return %0 : tensor<3x5x4xf32>
  // CHECK-LABEL: unary_einsum_reduce_sum_transpose
  // CHECK-DAG: %[[cst:.*]] = arith.constant dense<3> : tensor<1xi32>

// MinElement-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 @depth_to_space(%arg0: tensor<1x1x1x4xf32>) -> tensor<1x2x2x1xf32> {
  %0 = "tf.DepthToSpace"(%arg0) {block_size = 2: i64,  data_format = "NHWC"}: (tensor<1x1x1x4xf32>) -> tensor<1x2x2x1xf32>
  func.return %0 : tensor<1x2x2x1xf32>
// CHECK: %[[CUSTOM_0:.*]] = "tfl.custom"(%arg0) <{custom_code = "FlexDepthToSpace", custom_option = #tfl<const_bytes : "{{.*}}">}> : (tensor<1x1x1x4xf32>) -> tensor<1x2x2x1xf32>
// CHECK: return %[[CUSTOM_0]] : tensor<1x2x2x1xf32>
}

// CHECK:           %[[VAL_6:.*]] = "tfl.reshape"(%[[VAL_1]], %[[VAL_2]]) : (tensor<1xf32>, tensor<4xi32>) -> tensor<1x1x1x1xf32>
// CHECK:           %[[VAL_7:.*]] = "tfl.concatenation"(%[[VAL_5]], %[[VAL_6]]) <{axis = 3 : i32, fused_activation_function = "NONE"}> : (tensor<1x1x1x1xf32>, tensor<1x1x1x1xf32>) -> tensor<1x1x1x2xf32>