%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>

  %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

  %0 = "tfl.squeeze"(%arg0) {squeeze_dims = [0]}: (tensor<*xf32>) -> tensor<*xf32>
  func.return %0: tensor<*xf32>

// CHECK-LABEL: DontConvertSqueezeToReshape
// CHECK: %[[RESULT:.*]] = "tfl.squeeze"(%arg0)
// CHECK:  return %[[RESULT]]
}

func.func @ConvertSqueezeToReshapeOnMultiDynamicDims(%arg0: tensor<?x?xf32>) -> tensor<?x?xf32> {

// CHECK-DAG:         %[[CST_1:.*]] = "tf.Const"() <{value = dense<[1, 1504]> : tensor<2xi32>}> : () -> tensor<2xi32>
// CHECK:             %[[VAL_2:.*]] = "tf.Squeeze"(%[[VAL_1]]) <{squeeze_dims = [0]}> : (tensor<1x2xi32>) -> tensor<2xi32>
// CHECK:             %[[VAL_3:.*]] = "tf.Slice"(%[[ARG_0]], %[[VAL_2]], %[[CST_1]]) : (tensor<1x2944xi32>, tensor<2xi32>, tensor<2xi32>) -> tensor<1x1504xi32>

    SmallVector<int64_t, 4> begin_indices(value_rank, 0);
    auto end_indices = llvm::to_vector<4>(value_type.getShape());
    SmallVector<int64_t, 4> strides(value_rank, 1);

    // All HLO slice+squeeze results used to replace the original tf.Unpack op.
    SmallVector<Value, 4> results;
    results.reserve(op.getNumResults());

    for (int i = 0, end = op.getNumResults(); i < end; ++i) {