small cleanup

Author: lucidrains

native_sparse_attention_pytorch/native_sparse_attention.py

@@ -82,11 +82,16 @@ def round_up_mult(n, mult):
 def divisible_by(num, den):
     return (num % den) == 0
 
+# tensor helpers
+
 def pad_at_dim(t, pad, dim = -1, value = 0.):
     dims_from_right = (- dim - 1) if dim < 0 else (t.ndim - dim - 1)
     zeros = ((0, 0) * dims_from_right)
     return F.pad(t, (*zeros, *pad), value = value)
 
+def straight_through(t, target):
+    return t + (target - t).detach()
+
 # classes
 
 class SparseAttention(Module):
@@ -306,7 +311,7 @@ def forward(
             selected_importance_values, selected_block_indices = importance_scores.topk(num_selected, dim = -1)
 
             if self.use_diff_topk:
-                gates = selected_importance_values + (1. - selected_importance_values).detach()
+                gates = straight_through(selected_importance_values, 1.)
 
             fmask = selected_importance_values > 1e-10