add beartype decorator to RingRotaryEmbedding

lucidrains · Apr 11, 2024 · 5fd1e48 · 5fd1e48
1 parent aa7be57
commit 5fd1e48
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 1 deletion.
diff --git a/ring_attention_pytorch/ring_attention.py b/ring_attention_pytorch/ring_attention.py
@@ -40,6 +40,7 @@ def cast_tuple(t, length = 1):
 def divisible_by(num, den):
     return (num % den) == 0
 
+@beartype
 def default_attention(
     q: Tensor,
     k: Tensor,
@@ -80,6 +81,7 @@ def default_attention(
 # rotary embeddings with modifications to support striped attention
 
 class RingRotaryEmbedding(Module):
+    @beartype
     def __init__(
         self,
         dim,
@@ -105,6 +107,7 @@ def is_cuda(self):
         return self.inv_freq.is_cuda
 
     @autocast(enabled = False)
+    @beartype
     def forward(
         self,
         seq_len: int

diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
   name = 'ring-attention-pytorch',
   packages = find_packages(exclude=[]),
-  version = '0.3.20',
+  version = '0.3.21',
   license='MIT',
   description = 'Ring Attention - Pytorch',
   author = 'Phil Wang',