LogicalResult matchAndRewrite(OpTy op,
                                PatternRewriter &rewriter) const override {
    auto data_type = mlir::dyn_cast<RankedTensorType>(op.getData().getType());
    if (!data_type) return failure();
    int64_t data_rank = data_type.getRank();

    auto segment_ids_type =
        mlir::dyn_cast<RankedTensorType>(op.getSegmentIds().getType());