remove the get_at from inference

Author: lucidrains

native_sparse_attention_pytorch/__init__.py

@@ -1,7 +1,3 @@
 from native_sparse_attention_pytorch.native_sparse_attention import (
     SparseAttention
 )
-
-from native_sparse_attention_pytorch.triton_native_sparse_attention import (
-    native_sparse_attend
-)

native_sparse_attention_pytorch/native_sparse_attention.py

@@ -418,8 +418,14 @@ def forward_inference(
             sel_fk = rearrange(sel_fk, 'b h (w j) d -> b h w j d', j = self.selection_block_size)
             sel_fv = rearrange(sel_fv, 'b h (w j) d -> b h w j d', j = self.selection_block_size)
 
-            sel_fk = einx.get_at('b h [w] j d, b h 1 sel -> b h (sel j) d', sel_fk, sel_indices)
-            sel_fv = einx.get_at('b h [w] j d, b h 1 sel -> b h (sel j) d', sel_fv, sel_indices)
+            # get_at('b h [w] j d, b h 1 sel -> b h (sel j) d'
+
+            sel_indices = repeat(sel_indices, 'b h 1 sel -> b h sel j d', j = self.selection_block_size, d = sel_fk.shape[-1])
+
+            sel_fk = sel_fk.gather(2, sel_indices)
+            sel_fv = sel_fv.gather(2, sel_indices)
+
+            sel_fk, sel_fv = tuple(rearrange(t, 'b h sel j d -> b h (sel j) d') for t in (sel_fk, sel_fv))
 
             fmask = sel_scores > 1e-10
 

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "native-sparse-attention-pytorch"
-version = "0.0.64"
+version = "0.0.65"
 description = "Native Sparse Attention"
 authors = [
     { name = "Phil Wang", email = "[email protected]" }