if (!shift_ty || shift_ty.getRank() != 0) {
      return rewriter.notifyMatchFailure(
          op, "require the type of shift to be 0D tensor");
    }

    APInt val;
    if (!matchPattern(op.getAxis(), m_ConstantInt(&val))) {
      return rewriter.notifyMatchFailure(op, "require axis to be constant");
    }
    int axis = val.getSExtValue();