// Rewire output & get rid of the pack op.
  rewriter.replaceOp(pack_op, reshape_op.getResult());
  return success();
}

// ================== squared_difference ========================

LogicalResult SquaredDifference::matchAndRewrite(
    TFL::SquaredDifferenceOp squared_diff_op, PatternRewriter& rewriter) const {
  auto x = squared_diff_op.getLhs();
  auto y = squared_diff_op.getRhs();

  REDUCE_MIN = 89,
  FLOOR_DIV = 90,
  REDUCE_ANY = 91,
  SQUARE = 92,
  ZEROS_LIKE = 93,
  FILL = 94,
  FLOOR_MOD = 95,
  RANGE = 96,
  RESIZE_NEAREST_NEIGHBOR = 97,
  LEAKY_RELU = 98,
  SQUARED_DIFFERENCE = 99,
  MIRROR_PAD = 100,
  ABS = 101,
  SPLIT_V = 102,
  UNIQUE = 103,
  CEIL = 104,
  REVERSE_V2 = 105,
  ADD_N = 106,
  GATHER_ND = 107,
  COS = 108,
  WHERE = 109,
  RANK = 110,

  REDUCE_MIN = 89,
  FLOOR_DIV = 90,
  REDUCE_ANY = 91,
  SQUARE = 92,
  ZEROS_LIKE = 93,
  FILL = 94,
  FLOOR_MOD = 95,
  RANGE = 96,
  RESIZE_NEAREST_NEIGHBOR = 97,
  LEAKY_RELU = 98,
  SQUARED_DIFFERENCE = 99,
  MIRROR_PAD = 100,
  ABS = 101,
  SPLIT_V = 102,
  UNIQUE = 103,
  CEIL = 104,
  REVERSE_V2 = 105,
  ADD_N = 106,
  GATHER_ND = 107,
  COS = 108,
  WHERE = 109,
  RANK = 110,

  REDUCE_MIN = 89,
  FLOOR_DIV = 90,
  REDUCE_ANY = 91,
  SQUARE = 92,
  ZEROS_LIKE = 93,
  FILL = 94,
  FLOOR_MOD = 95,
  RANGE = 96,
  RESIZE_NEAREST_NEIGHBOR = 97,
  LEAKY_RELU = 98,
  SQUARED_DIFFERENCE = 99,
  MIRROR_PAD = 100,
  ABS = 101,
  SPLIT_V = 102,
  UNIQUE = 103,
  CEIL = 104,
  REVERSE_V2 = 105,
  ADD_N = 106,
  GATHER_ND = 107,
  COS = 108,
  WHERE = 109,
  RANK = 110,

// In above calculation, they are replaced by new values. These new mean and
// variance are calculated as following:
// new_mean = mean(x, axis=[0, 1, 2])
// new_variance = mean(squared_difference(x, new_mean), axis=[0, 1, 2])
//
// The DDR rule for the is_training equals true case is as following:
// def : Pattern<
//     (TF_FusedBatchNormV3Op:$root
//         $x, $scale, $offset, $mean, $variance,

  %0 = "tfl.squared_difference"(%arg0, %arg1) : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
  %1 = "tfl.relu"(%0) : (tensor<1xf32>) -> tensor<1xf32>
  func.return %1: tensor<1xf32>

// CHECK-LABEL: squaredDifferenceReluRemoveRelu
// CHECK:  %[[RESULT:.*]] = tfl.squared_difference %arg0, %arg1 : tensor<1xf32>
// CHECK:  return %[[RESULT]]
}

  using OpRewritePattern<TFL::PackOp>::OpRewritePattern;

  LogicalResult matchAndRewrite(TFL::PackOp pack_op,
                                PatternRewriter& rewriter) const override;
};

struct SquaredDifference : public OpRewritePattern<TFL::SquaredDifferenceOp> {
  using OpRewritePattern<TFL::SquaredDifferenceOp>::OpRewritePattern;

  LogicalResult matchAndRewrite(TFL::SquaredDifferenceOp squared_diff_op,

^bb0(%arg0: tensor<1xf32>, %arg1: tensor<1xf32>):
  %0 = "tf.SquaredDifference"(%arg0, %arg1) : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
  %1 = "tf.Relu6"(%0) : (tensor<1xf32>) -> tensor<1xf32>
  func.return %1: tensor<1xf32>

// CHECK-LABEL: squaredDifferenceRelu
// CHECK:  tfl.squared_difference %arg0, %arg1 : tensor<1xf32>
// CHECK:  %1 = "tfl.relu6"(%0) : (tensor<1xf32>) -> tensor<1xf32>

    }
    void $cppClass::print(OpAsmPrinter &p) {
      return printOneResultOp(getOperation(), p);
    }
  }];

  let hasOptions = 1;
}

def TFL_SquaredDifferenceOp : TFL_Op<"squared_difference", [
    TFL_OperandsHaveSameShapesOrBroadcastableShape<[0, 1], 4>,
    BinaryOpSameElementTypeConstraint,
    TFL_SameFirstOperandAndFirstResultElementType,
    ResultsBroadcastableShape,

  %0 = "tfl.squared_difference"(%arg0, %arg1) : (tensor<1x80x128x!quant.uniform<i8:f32, 0.089839041233062744:10>>, tensor<1x80x128x!quant.uniform<i8:f32, 0.0019308560295030475:-6>>) -> tensor<1x80x128x!quant.uniform<i8:f32, 0.60070550441741943:-128>>