%0 = "tf.Const"() {value = dense<0.0> : tensor<1x2xf32>} : () -> tensor<1x2xf32>

  // Operand and identity shapes are broadcastable. However, we cannot fold
  // because the operand does not match the result shape.
  %1 = "tf.AddV2"(%arg0, %0) : (tensor<4x1xf32>, tensor<1x2xf32>) -> tensor<4x2xf32>

      time_major = false} : (
        tensor<1x2x3xf32>,
        tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>,
        tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>,
        none, none, none,
        tensor<3xf32>, tensor<3xf32>, tensor<3xf32>, tensor<3xf32>,
        none, none,
        tensor<1x3xf32>, tensor<1x3xf32>,
        none, none, none, none) -> tensor<1x2x3xf32>

// CHECK: return %[[GATHER]] : tensor<1x300x10xf32>

// -----

module {
  func.func @propagate_xlagather(%arg0: tensor<10x2xi32>) -> tensor<1x300x10xf32> {
    %cst = "tf.Const"() {value = dense<[1, 1, 300]> : tensor<3xi64>} : () -> tensor<3xi64>

  %1 = "tf.ReadVariableOp"(%0) {device = "/CPU:0"} : (tensor<!tf_type.resource<tensor<1x1xf32>>>) -> tensor<1x1xf32>
  %2:2 = "tf.Split"(%cst, %arg0) {device = "/CPU:0"} : (tensor<i32>,  tensor<1x4xf32>) -> (tensor<1x2xf32>, tensor<1x2xf32>)
  %3:2 = "tf.Split"(%cst, %2#0) {device = "/CPU:0"} : (tensor<i32>,  tensor<1x2xf32>) -> (tensor<1x1xf32>, tensor<1x1xf32>)

// CHECK: %0 = "tfl.pseudo_const"() <{value = dense<2> : tensor<2xi32>}> : () -> tensor<2xi32>
// CHECK: %1 = "tfl.reshape"(%arg0, %0) : (tensor<1x4xf32>, tensor<2xi32>) -> tensor<2x2xf32>

{
  "version": 3,
  "operator_codes": [
    {
      "builtin_code": "RESHAPE"
    }
  ],
  "subgraphs": [
    {
      "tensors": [
        {
          "shape": [1, 4],

// `tfl.concatenation`.

func.func @concatenate_float(%arg0: tensor<3x2xf32>, %arg1: tensor<1x2xf32>) -> tensor<4x2xf32> {
  %0 = "stablehlo.concatenate"(%arg0, %arg1) {dimension = 0 : i64} : (tensor<3x2xf32>, tensor<1x2xf32>) -> tensor<4x2xf32>
  return %0 : tensor<4x2xf32>
}
// CHECK-LABEL: concatenate_float
// CHECK-NOT: tfl.concatenation
// CHECK: stablehlo.concatenate

// -----

// CHECK-DAG: %[[input_18:.*]] = "quantfork.stats"({{.*}}) <{layerStats = dense<[-8.000000e-01, 1.600000e+00]> : tensor<2xf32>}> : (tensor<1x4xf32>) -> tensor<1x4xf32>
// CHECK-DAG: %[[input_19:.*]] = "quantfork.stats"({{.*}}) <{layerStats = dense<[-2.000000e+00, 4.000000e+00]> : tensor<2xf32>}> : (tensor<1x2xf32>) -> tensor<1x2xf32>

// CHECK: "tfl.unidirectional_sequence_lstm"({{.*}}, %[[input_18]], %[[input_19]], %{{[0-9]+}}, %{{[0-9]+}}, %{{[0-9]+}}, %{{[0-9]+}})

// RUN:   -split-input-file | FileCheck %s --check-prefix=STATIC-RANGE-PTQ-TO-COMPUTE-HEAVY

// STATIC-RANGE-PTQ-TO-COMPUTE-HEAVY: @main
func.func @main(%arg0: tensor<1x2xf32>) -> tensor<1x2xf32> {
  %0 = stablehlo.add %arg0, %arg0 : tensor<1x2xf32>
  return %0 : tensor<1x2xf32>
}
// Tests that `composite_add_fn_1` does not quantize when quantizing
// only compute-heavy ops.

  %9 = "mhlo.compare"(%0, %8) {comparison_direction = #mhlo<comparison_direction NE>} : (tensor<10x8xf32>, tensor<10x8xf32>) -> tensor<10x8xi1>
  %10 = mhlo.and %7, %9 : tensor<10x8xi1>
  %11 = mhlo.subtract %arg0, %6 : tensor<10x8xf32>
  %12 = mhlo.divide %11, %5 : tensor<10x8xf32>
  %13 = mhlo.add %12, %3 : tensor<10x8xf32>

// CHECK-SAME: (%[[ARG_0:[a-z0-9]*]]: tensor<!tf_type.resource<tensor<10x3xf32>>>, %[[ARG_1:[a-z0-9]*]]: tensor<!tf_type.resource<tensor<10x3xf32>>>, %[[ARG_2:[a-z0-9]*]]: tensor<!tf_type.resource<tensor<10x3xf32>>>, %[[ARG_3:[a-z0-9]*]]: tensor<!tf_type.resource<tensor<10x3xf32>>>)
!rtype = tensor<!tf_type.resource<tensor<10x3xf32>>>