Operation* consumer_op = maybeConsumer.second;

  TF::ExpandDimsOp expand_op;
  TF::SqueezeOp squeeze_op;
  int64_t expand_axis = -1;
  // Expand + Squeeze op.
  if (llvm::isa<TF::ExpandDimsOp>(producer_op)) {
    if (!llvm::isa<TF::SqueezeOp>(consumer_op)) {
      // Expand/Squeeze op must come in pair.
      return rewriter.notifyMatchFailure(
          op, "ExpandDimsOp and SqueezeOp should come in pair");
    }

  // has one result operand.
  return isa<TF::IdentityOp, TF::CastOp, TF::ReshapeOp, TF::XlaShardingOp,
             TF::GatherOp, TF::GatherV2Op, TF::XlaGatherOp, TF::ExpandDimsOp,
             TF::SqueezeOp, TF::TransposeOp>(op);
}

bool IsOpWithQuantizableTrait(Operation* op) {
  // Supported quantizable ops.
  return isa<TF::XlaConvV2Op, TF::XlaDotV2Op, TF::MatMulOp, TF::Conv2DOp,

namespace quant {
namespace {

// This pass duplicates constants that affect or determine the shape of a tensor
// after being used in a computation for some op. Some specific operands of TF
// ops (like the `dim` argument for `TF::ExpandDimsOp`) determine the shape of
// the resulting tensor. If these operands are constants, they are duplicated
// and replace the shape-determining operands. Each duplicated constant will

  Value scalar_zero = CreateI32SplatConst(loc, rewriter, {}, 0);
  RankedTensorType vector_type =
      tensorflow::GetTypeFromTFTensorShape({1}, shape_dtype);
  auto expanded_start_index = rewriter->create<TF::ExpandDimsOp>(
      loc, vector_type, start_index, scalar_zero);
  auto start_position = rewriter->create<TF::ConcatOp>(
      loc, position_type, scalar_zero,
      ArrayRef<Value>({expanded_start_index, partial_start_position}));

  //
  // TODO(chky): Find a more scalable way to figure out cheap ops.
  if (llvm::isa<mlir::TF::ShapeOp, mlir::TF::StridedSliceOp,
                mlir::TF::ReshapeOp, mlir::TF::ExpandDimsOp>(op)) {
    cost_map_[op] = kDefaultCheapCost;
    return;
  }

  // For other ops, use the sum of input sizes as its cost.
  int64_t cost = kDefaultCheapCost;
  for (auto operand : op->getOperands()) {

    TypeID::get<TF::DepthwiseConv2dNativeOp>(),
    TypeID::get<TF::DivOp>(),
    TypeID::get<TF::DynamicStitchOp>(),
    TypeID::get<TF::_EagerConstOp>(),
    TypeID::get<TF::EmptyOp>(),
    TypeID::get<TF::ExpandDimsOp>(),
    TypeID::get<TF::FakeQuantWithMinMaxVarsOp>(),
    TypeID::get<TF::FillOp>(),
    TypeID::get<TF::FusedBatchNormOp>(),
    TypeID::get<TF::FusedBatchNormGradOp>(),
    TypeID::get<TF::FusedBatchNormGradV2Op>(),

        r++;
      }
    }

    auto expaned_input_type = tensorflow::GetTypeFromTFTensorShape(
        expanded_input_shape, input_type.getElementType());
    TF::ExpandDimsOp expanded_input = rewriter.create<TF::ExpandDimsOp>(
        rfft_op.getLoc(), expaned_input_type, input, minus_two->getResult());

    // Expanded fft_len.
    auto one_attr = mlir::DenseIntElementsAttr::get(one_ele_type, {1});

 public:
  explicit LowerPackOp(MLIRContext *context)
      : RewritePattern(
            PackOp::getOperationName(), 1, context,
            {ConstOp::getOperationName(), ConcatV2Op::getOperationName(),
             ExpandDimsOp::getOperationName()}) {}

  LogicalResult matchAndRewrite(Operation *src_op,
                                PatternRewriter &rewriter) const override {
    auto op = cast<PackOp>(src_op);

                                             MLIRContext* context) {
  results.add(flipComatibleShapeError<NotEqualOp>);
}

//===----------------------------------------------------------------------===//
// ExpandDimsOp
//===----------------------------------------------------------------------===//

Type InferExpandDimsOpType(Value input, Value dim) {
  Type element_ty = mlir::cast<TensorType>(input.getType()).getElementType();

    return success();
  }
};

class ConvertDynamicExpandDimsOp : public OpRewritePattern<TF::ExpandDimsOp> {
 public:
  using OpRewritePattern::OpRewritePattern;

  LogicalResult matchAndRewrite(TF::ExpandDimsOp op,
                                PatternRewriter &rewriter) const override {
    auto input = op.getInput();