}

TEST_F(ArrayGradTest, SqueezeGrad) {
  TensorShape x_shape({1, 5, 1, 2});
  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
  TensorShape y_shape({5, 2});
  auto y = Squeeze(scope_, x);
  RunTest(x, x_shape, y, y_shape);
}

TEST_F(ArrayGradTest, TransposeGrad) {
  TensorShape x_shape({5, 2});
  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));

    %12 = "tf.Cast"(%arg1) {Truncate = false} : (tensor<2x1xf32>) -> tensor<2x1xi64>
    %13 = "tf.Reshape"(%12, %9) : (tensor<2x1xi64>, tensor<1xi32>) -> tensor<2xi64>
    %14 = "tf.Squeeze"(%arg1) {squeeze_dims = [-1]} : (tensor<2x1xf32>) -> tensor<2xf32>
    // CHECK: "tf.Conv2D"
    // CHECK-SAME: strides = [1, 1, 1, 1]

  %267 = "tf.Reshape"(%266, %5) {device = ""} : (tensor<?x?xi1>, tensor<1xi32>) -> tensor<?xi1>
  %268 = "tf.Where"(%267) {device = ""} : (tensor<?xi1>) -> tensor<?x1xi64>
  %269 = "tf.Squeeze"(%268) {device = "", squeeze_dims = [1]} : (tensor<?x1xi64>) -> tensor<?xi64>
  %270 = "tf.GatherV2"(%264, %269, %14) {batch_dims = 0 : i64, device = ""} : (tensor<?xi64>, tensor<?xi64>, tensor<i32>) -> tensor<?xi64>

  auto float_graph = readonly_model_->subgraphs()->Get(0);
  // The original model reshape->custom->custom->squeeze.
  ASSERT_THAT(*float_graph->operators(), SizeIs(4));
  // The resulting model should be:
  // reshape->dequantize->custom->custom->quantize->squeeze.
  ASSERT_THAT(subgraph->operators, SizeIs(6));
  const std::vector<BuiltinOperator> op_codes = {

func.func @PadStridedSliceNewAxisMask2(%arg0: tensor<4x64x64x1xf32>) -> tensor<1x4x64x64xf32> {
  %cst = arith.constant dense<0> : tensor<3xi32>
  %cst_0 = arith.constant dense<1> : tensor<3xi32>
  %0 = "tf.Squeeze"(%arg0) {T = f32, _output_shapes = ["tfshape$dim { size: 4 } dim { size: 64 } dim { size: 64 }"], device = "", squeeze_dims = []} : (tensor<4x64x64x1xf32>) -> tensor<4x64x64xf32>

  PAD = 34,
  UNIDIRECTIONAL_SEQUENCE_RNN = 35,
  GATHER = 36,
  BATCH_TO_SPACE_ND = 37,
  SPACE_TO_BATCH_ND = 38,
  TRANSPOSE = 39,
  MEAN = 40,
  SUB = 41,
  DIV = 42,
  SQUEEZE = 43,
  UNIDIRECTIONAL_SEQUENCE_LSTM = 44,
  STRIDED_SLICE = 45,
  BIDIRECTIONAL_SEQUENCE_RNN = 46,
  EXP = 47,
  TOPK_V2 = 48,
  SPLIT = 49,
  LOG_SOFTMAX = 50,

  %0 = "tf.Bitcast"(%arg0) : (tensor<2xf32>) -> tensor<2x2xf16>
  func.return %0 : tensor<2x2xf16>
}

// -----

// CHECK-LABEL: squeeze
func.func @squeeze(%arg0: tensor<1x1x10xf32>) -> tensor<1x10xf32> {
  // CHECK: mhlo.reshape
  %0 = "tf.Squeeze"(%arg0) : (tensor<1x1x10xf32>) -> tensor<1x10xf32>
  func.return %0 : tensor<1x10xf32>
}

// -----

// CHECK-LABEL: squeeze_ranked

  PAD = 34,
  UNIDIRECTIONAL_SEQUENCE_RNN = 35,
  GATHER = 36,
  BATCH_TO_SPACE_ND = 37,
  SPACE_TO_BATCH_ND = 38,
  TRANSPOSE = 39,
  MEAN = 40,
  SUB = 41,
  DIV = 42,
  SQUEEZE = 43,
  UNIDIRECTIONAL_SEQUENCE_LSTM = 44,
  STRIDED_SLICE = 45,
  BIDIRECTIONAL_SEQUENCE_RNN = 46,
  EXP = 47,
  TOPK_V2 = 48,
  SPLIT = 49,
  LOG_SOFTMAX = 50,

  const auto& expected_params = collected_quant_params[0];
  for (int i = 1; i < collected_quant_params.size(); ++i) {
    const auto& compared_params = collected_quant_params[i];
    // For some ops (such as Transpose or Squeeze), the quantized axis might not
    // be the same, this function only verifies the scale and zero point in
    // that case. The quantized axis should be verified in their own verifier
    // method.

This operation is related to `squeeze()`, which removes dimensions of
size 1.
  }];

  // TODO: Restriction on dim's size and valid range are not modeled here.
  let arguments = (ins AnyTensor:$input, TFL_I32OrI64Tensor:$dim);

  let results = (outs AnyTensor:$output);

  let hasOptions = 1;
}

def TFL_SqueezeOp: TFL_Op<"squeeze", [Pure,