Is there a version of the rotary embedding with unpadded inputs

Thibault712 · July 11, 2023, 7:40pm

HazyResearch/flash-attention/blob/main/flash_attn/layers/rotary.py

# Copyright (c) 2023, Tri Dao.

from typing import Tuple
import math

import torch

from einops import rearrange, repeat

import rotary_emb


def rotate_half(x, interleaved=False):
    if not interleaved:
        x1, x2 = x.chunk(2, dim=-1)
        return torch.cat((-x2, x1), dim=-1)
    else:
        x1, x2 = x[..., ::2], x[..., 1::2]
        return rearrange(torch.stack((-x2, x1), dim=-1), '... d two -> ... (d two)', two=2)

This file has been truncated. show original

I can see that this one working with flash attention is only for padded inputs.