unsigned num_arguments = main_fn.getNumArguments();
  for (unsigned arg_index = 0; arg_index < num_arguments; ++arg_index) {
    if (auto aliasing_output = main_fn.getArgAttrOfType<mlir::IntegerAttr>(
            arg_index, kAliasingAttr))
      output_to_input_alias[aliasing_output.getInt()] = arg_index;
  }

  if (output_to_input_alias.empty()) return absl::OkStatus();

  using OpRewritePattern<GetElementOp>::OpRewritePattern;

 public:
  LogicalResult matchAndRewrite(GetElementOp ge_op,
                                PatternRewriter &rewriter) const override {
    IntegerAttr index;
    if (!matchPattern(ge_op.getIndex(), m_Constant(&index))) {
      return failure();
    }
    auto preceding_build_list = llvm::dyn_cast_or_null<BuildListOp>(
        ge_op.getTensorList().getDefiningOp());

      ArrayAttr targs_attr = ArrayAttr::get(context, targs_values);
      attrs.push_back(
          NamedAttribute(StringAttr::get(context, "TArgs"), targs_attr));

      auto num_args_attr = IntegerAttr::get(IntegerType::get(context, 64), 1);
      attrs.push_back(
          NamedAttribute(StringAttr::get(context, "num_args"), num_args_attr));

      // Fused conv operands are input, filter, args and host args. Here, bias

    SmallVector<int32_t, 4> perm;
    SmallVector<int64_t, 4> output_shape;
    for (int i = 0; i < num_dimensions; ++i) {
      perm.push_back(perm_tensor.getValues<IntegerAttr>()[i].getInt());
      output_shape.push_back(input_shape[perm[i]]);

      // Check that the derived output shape matches the static shape.
      assert(!output_type.hasStaticShape() ||

    if (auto float_attr = mlir::dyn_cast_or_null<FloatAttr>(attr.second)) {
      fbb.Float(attr.first.data(), float_attr.getValue().convertToFloat());
    } else if (auto int_attr =
                   mlir::dyn_cast_or_null<IntegerAttr>(attr.second)) {
      fbb.Int(attr.first.data(), int_attr.getInt());
    } else if (auto bool_attr = mlir::dyn_cast_or_null<BoolAttr>(attr.second)) {
      fbb.Bool(attr.first.data(), bool_attr.getValue());

  // If the validation fails, the pass should stop immediately.
  if (!IsLegalQuantSpecs(func)) {
    return true;
  }

  OpBuilder builder(func);
  bool is_signed = quant_specs_.IsSignedInferenceType();
  IntegerAttr num_bits =
      builder.getI32IntegerAttr(quant_specs_.GetQuantizationTypeWidth());
  BoolAttr narrow_range = builder.getBoolAttr(false);

  auto add_quantize_op = [&](Location loc, Type input_type, Block* block,

  // If the validation fails, the pass should stop immediately.
  if (!IsLegalQuantSpecs(func)) {
    return true;
  }

  OpBuilder builder(func);
  bool is_signed = quant_specs_.IsSignedInferenceType();
  IntegerAttr num_bits =
      builder.getI32IntegerAttr(quant_specs_.GetQuantizationTypeWidth());
  BoolAttr narrow_range = builder.getBoolAttr(false);

  auto add_quantize_op = [&](Location loc, Type input_type, Block* block,

  DenseIntElementsAttr axis_attr;
  if (!matchPattern(op.getAxis(), m_Constant(&axis_attr))) return failure();
  if (axis_attr.getNumElements() != 1) return failure();
  int64_t axis =
      axis_attr.getSplatValue<IntegerAttr>().getValue().getSExtValue();
  // TODO(ezhulenev): Compute axis from rank. e.g. It might be common to concat
  // on the channels dim for NCHW layout as axis=-2.
  if (axis < 0) return failure();

      stride_w = CastI64ToI32((*window_strides)[2]).value();
    }
    auto stride_h_attr = IntegerAttr::get(rewriter.getI32Type(), stride_h);
    auto stride_w_attr = IntegerAttr::get(rewriter.getI32Type(), stride_w);

    ArrayRef<int64_t> window_dims = op.getWindowDimensions();
    auto window_w_attr = IntegerAttr::get(rewriter.getI32Type(),
                                          CastI64ToI32(window_dims[2]).value());

    if (squeeze_dims.size() != 1) {
      return rewriter.notifyMatchFailure(
          squeeze_op, "squeeze dims should have exactly 1 dimension specified");
    }
    int64_t squeeze_axis = mlir::cast<IntegerAttr>(squeeze_dims[0]).getInt();
    if (squeeze_axis < 0) {
      // Always squeeze 4D input to 3D input.
      squeeze_axis += 4;
    }
    if (squeeze_axis != expand_axis) {
      return rewriter.notifyMatchFailure(