%0 = mhlo.composite "aten.avg_pool2d.default" %arg0 {composite_attributes = {ceil_mode = false, count_include_pad = true, divisor_override = "py_None", kernel_size = dense<3> : tensor<2xi64>, padding = dense<0> : tensor<2xi64>, stride = dense<1> : tensor<2xi64>}, decomposition = @XlaCallModule_aten.avg_pool2d.default.impl_0} : (tensor<1x3x6x6xf32>) -> tensor<1x3x4x4xf32>
  %1 = "tf.Identity"(%0) {device = ""} : (tensor<1x3x4x4xf32>) -> tensor<*xf32>

  // CHECK-NEXT: %12 = "mhlo.pad"(%10, %11) <{edge_padding_high = dense<[3, 4]> : tensor<2xi64>, edge_padding_low = dense<0> : tensor<2xi64>, interior_padding = dense<0> : tensor<2xi64>}> : (tensor<3x4xi32>, tensor<i32>) -> tensor<6x8xi32>
  // CHECK-NEXT: %13 = "mhlo.slice"(%arg1) <{limit_indices = dense<1> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>}> : (tensor<2xi32>) -> tensor<1xi32>

// CHECK-NEXT:      %16 = "tfl.pseudo_const"() <{value = dense<[0, 1]> : tensor<2xi64>}> : () -> tensor<2xi64>
// CHECK-NEXT:      %17 = "tfl.gather"(%14, %16) <{axis = 0 : i32, batch_dims = 0 : i32}> : (tensor<3xi32>, tensor<2xi64>) -> tensor<2xi32>
// CHECK-NEXT:      %18 = "tfl.pseudo_const"() <{value = dense<1> : tensor<1xi64>}> : () -> tensor<1xi64>

  %cst_0 = arith.constant dense<[2, -1]> : tensor<2xi64>
  %0 = "tf.FakeQuantWithMinMaxVars"(%arg0, %arg1, %arg2) {num_bits = 5, narrow_range = false} : (tensor<1x2xf32>, tensor<f32>, tensor<f32>) -> tensor<1x2xf32>
  %1 = "tf.Reshape"(%0, %cst_0) : (tensor<1x2xf32>, tensor<2xi64>) -> tensor<2x1xf32>
  %2 = "tf.Reshape"(%0, %cst_0) : (tensor<1x2xf32>, tensor<2xi64>) -> tensor<2x1xf32>
  func.return %1, %2 : tensor<2x1xf32>, tensor<2x1xf32>

  %cst_0 = arith.constant dense<[2, -1]> : tensor<2xi64>
  %0 = "tf.FakeQuantWithMinMaxVars"(%arg0, %arg1, %arg2) {num_bits = 3, narrow_range = false} : (tensor<1x2xf32>, tensor<f32>, tensor<f32>) -> tensor<1x2xf32>
  %1 = "tf.Reshape"(%0, %cst_0) : (tensor<1x2xf32>, tensor<2xi64>) -> tensor<2x1xf32>
  %2 = "tf.Reshape"(%0, %cst_0) : (tensor<1x2xf32>, tensor<2xi64>) -> tensor<2x1xf32>
  func.return %1, %2 : tensor<2x1xf32>, tensor<2x1xf32>

    // CHECK: "mhlo.pad"(%[[ARG0]], %[[ARG1]])
    // CHECK-SAME-DAG: edge_padding_high = dense<[1, 2]> : tensor<2xi64>
    // CHECK-SAME-DAG: edge_padding_low = dense<[2, 1]> : tensor<2xi64>
    // CHECK-SAME-DAG: interior_padding = dense<[1, 0]> : tensor<2xi64>

    %0 = mhlo.constant dense<[2, 1]> : tensor<2xi32>
    %1 = mhlo.constant dense<[1, 2]> : tensor<2xi32>

  // CHECK:  return %[[CST]]
}

// CHECK-LABEL: @ShapeOpI64
func.func @ShapeOpI64(%arg0 : tensor<576x72xf32>) -> tensor<2xi64> {
  %0 = "tfl.shape"(%arg0) : (tensor<576x72xf32>) -> tensor<2xi64>
  func.return %0 : tensor<2xi64>
  // CHECK: %[[CST:.*]] = arith.constant dense<[576, 72]> : tensor<2xi64>
  // CHECK:  return %[[CST]]
}

// CHECK-LABEL: @ConstFoldStridedSlice

    %cst = arith.constant dense<[2, 100]> : tensor<2xi64>
    %0 = "tfl.concatenation"(%arg0, %arg1) {axis = 0 : i32, fused_activation_function = "NONE"} : (tensor<100xf32>, tensor<100xf32>) -> tensor<200xf32>
    %1 = "tfl.reshape"(%0, %cst) : (tensor<200xf32>, tensor<2xi64>) -> tensor<2x100xf32>
    func.return %1 : tensor<2x100xf32>
  }

  // CHECK-NEXT:   %[[LHS_BCAST:.+]] = "mhlo.dynamic_broadcast_in_dim"(%arg0, %[[RESULT_EXTENTS]]) <{broadcast_dimensions = dense<1> : tensor<1xi64>}>
  // CHECK-NEXT:   %[[RHS_BCAST:.+]] = "mhlo.dynamic_broadcast_in_dim"(%arg1, %[[RESULT_EXTENTS]]) <{broadcast_dimensions = dense<[0, 1]> : tensor<2xi64>}>
  // CHECK-NEXT:   %[[RESULT:.+]] = mhlo.add %[[LHS_BCAST]], %[[RHS_BCAST]] : tensor<?x?xi32>
  // CHECK-NEXT:   shape.assuming_yield %[[RESULT]]

      while (auto cast_op =
                 llvm::dyn_cast_or_null<CastOp>(value.getDefiningOp())) {
        // Consider cast compatibility in case
        //    %cast = "tf.Cast"(%0) : (tensor<2xi64>) -> tensor<2xf32>
        // is skipped.
        if (cast_op.getSrcT() != cast_op.getDstT()) {
          break;
        }
        value = cast_op.getOperand();
      }
      return value;
    };