if the chunk is to be all masked out causally, skip summarizing the block entirely, for memory efficient attention

Author: lucidrains

memory_efficient_attention_pytorch/memory_efficient_attention.py

@@ -102,6 +102,10 @@ def memory_efficient_attention(
 
             causal_mask_chunk = causal_mask_chunks[q_index][k_index] if causal else None
 
+            if exists(causal_mask_chunk) and torch.all(causal_mask_chunk):
+                # if chunk is to be all masked out causally, skip
+                continue
+
             exp_weight_chunk, weighted_value_chunk, weight_max_chunk = checkpointed_summarize_qkv_chunk(
                 q_chunk,
                 k_chunk,

setup.py

@@ -3,7 +3,7 @@
 setup(
   name = 'memory-efficient-attention-pytorch',
   packages = find_packages(exclude=[]),
-  version = '0.0.5',
+  version = '0.0.6',
   license='MIT',
   description = 'Memory Efficient Attention - Pytorch',
   author = 'Phil Wang',