Skip to content

Commit

Permalink
cleanup
Browse files Browse the repository at this point in the history
  • Loading branch information
lucidrains committed Dec 1, 2023
1 parent 204686d commit 02a5d07
Showing 2 changed files with 7 additions and 7 deletions.
6 changes: 5 additions & 1 deletion rotary_embedding_torch/__init__.py
Original file line number Diff line number Diff line change
@@ -1 +1,5 @@
from rotary_embedding_torch.rotary_embedding_torch import apply_rotary_emb, RotaryEmbedding, broadcat, apply_learned_rotations
from rotary_embedding_torch.rotary_embedding_torch import (
apply_rotary_emb,
RotaryEmbedding,
apply_learned_rotations
)
8 changes: 2 additions & 6 deletions rotary_embedding_torch/rotary_embedding_torch.py
Original file line number Diff line number Diff line change
@@ -18,10 +18,6 @@ def exists(val):
def default(val, d):
return val if exists(val) else d

def broadcat(tensors, dim = -1):
broadcasted_tensors = broadcast_tensors(*tensors)
return torch.cat(broadcasted_tensors, dim = dim)

# rotary embedding helper functions

def rotate_half(x):
@@ -225,8 +221,8 @@ def get_axial_freqs(self, *dims):
new_axis_slice = (Ellipsis, *all_axis, Colon)
all_freqs.append(freqs[new_axis_slice])

freqs = broadcat(all_freqs, dim = -1)
return freqs
all_freqs = broadcast_tensors(*all_freqs)
return torch.cat(all_freqs, dim = -1)

@autocast(enabled = False)
def forward(self, t, cache_key = None):

0 comments on commit 02a5d07

Please sign in to comment.