//   %cst = arith.constant dense<1.0> : tensor<16x16x4xf32>
//   %0 = "tfl.conv_2d"...
//   %1 = "tfl.add"(%0, %cst) : (tensor<16x16x4xf32>, tensor<16x16x4xf32>)
// After this optimization:
//   %cst = arith.constant dense<1.0> : tensor<f32>
//   %0 = "tfl.conv_2d"...
//   %1 = "tfl.add"(%0, %cst) : (tensor<16x16x4xf32>, tensor<f32>)

// input shape.

// CHECK-LABEL: reduce_window_max_activation_transpose
func.func @reduce_window_max_activation_transpose(%arg0: tensor<1x16x16x4xf32>) -> tensor<1x4x8x8xf32> {
  %0 = stablehlo.constant dense<0xFF800000> : tensor<f32>  // -inf
  %1 = stablehlo.transpose %arg0, dims = [0, 3, 1, 2] : (tensor<1x16x16x4xf32>) -> tensor<1x4x16x16xf32>
  %2 = "stablehlo.reduce_window"(%1, %0) ({
  ^bb0(%arg1: tensor<f32>, %arg2: tensor<f32>):

// DISABLE-ALL-DOT-GENERAL: @main
func.func @main(%arg0: tensor<1x1x167xf32>) -> tensor<1x1x64xf32> {
  %0 = stablehlo.constant dense<2.000000e+00> : tensor<167x64xf32>
  %1 = stablehlo.dot_general %arg0, %0, contracting_dims = [2] x [0], precision = [DEFAULT, DEFAULT] : (tensor<1x1x167xf32>, tensor<167x64xf32>) -> tensor<1x1x64xf32>
  return %1 : tensor<1x1x64xf32>
}

// DISABLE-ALL-DOT-GENERAL: %[[CONST:.+]] = stablehlo.constant dense<2.000000e+00>

  // CHECK: return [[RI]]
  func.return %ri : tensor<!tf_type.resource<tensor<16x16x16xf32>>>
}

// -----

// CHECK-LABEL:func @partitioned_output_3d
// CHECK-SAME: ([[ARG:%.*]]: tensor<!tf_type.resource<tensor<16x16x16xf32>>>)
func.func @partitioned_output_3d(%arg: tensor<!tf_type.resource<tensor<16x16x16xf32>>>) -> tensor<!tf_type.resource<tensor<16x8x16xf32>>> {

    %19 = stablehlo.convert %18 : (tensor<8x16x4xi32>) -> tensor<8x16x4xf32>
    %20 = stablehlo.broadcast_in_dim %9, dims = [0, 1, 2] : (tensor<1x1x1xf32>) -> tensor<8x16x4xf32>
    %21 = stablehlo.multiply %19, %20 : tensor<8x16x4xf32>  // * s1 s2
    %22 = call @uniform_quantize_1(%21, %7, %8) : (tensor<8x16x4xf32>, tensor<1x1x1xf32>, tensor<1x1x1xi8>) -> tensor<8x16x4xi8>

  %5 = stablehlo.broadcast_in_dim %1, dims = [2] : (tensor<64xf32>) -> tensor<1x1x64xf32>
  %6 = stablehlo.add %4, %5 : tensor<1x1x64xf32>
  %7 = stablehlo.clamp %2, %6, %3 : tensor<1x1x64xf32>
  func.return %7: tensor<1x1x64xf32>
}
// CHECK: %[[CONST_0:.*]] = stablehlo.constant dense<2.000000e+00>
// CHECK: %[[CONST_1:.*]] = stablehlo.constant dense<2.000000e+00>

      return %0 : tensor<1x1x4xf32>
    }
  )mlir";

  OwningOpRef<ModuleOp> module_op =
      ParseModuleOpString(kModuleXlaCallModuleOpWithDefaultQuantizationMethod);
  ASSERT_TRUE(module_op);

  %1 = "tfl.arg_max"(%0, %cst) : (tensor<16x1024xf32>, tensor<1xi32>) -> tensor<16xi32>
  func.return %1 : tensor<16xi32>
  // CHECK-DAG: %[[CST:.*]] = arith.constant dense<-1> : tensor<1xi32>
  // CHECK: %[[SOFTMAX:.*]] = "tfl.softmax"(%arg0) <{beta = -1.000000e+00 : f32}> : (tensor<16x1024xf32>) -> tensor<16x1024xf32>

      return %0 : tensor<1x1x4xf32>
    }
  )mlir";

  const OwningOpRef<ModuleOp> module_op =
      ParseModuleOpString(kXlaCallModuleOpWithQuantizationMethodAttr);
  ASSERT_TRUE(module_op);

  func::FuncOp main_fn = FindMainFuncOp(*module_op);

  %split_dim = arith.constant dense<0> : tensor<2x2xi32>
  // expected-error @+1 {{'tfl.split' op failed to verify that operand #0 is an 0-d tensor or 1-d tensor w/ 1 element}}
  %0 = "tfl.split"(%split_dim, %arg0) {num_splits = 1 : i32} : (tensor<2x2xi32>, tensor<16x4x4xf32>) -> tensor<16x4x4xf32>
  func.return %0 : tensor<16x4x4xf32>
}

// -----