stack-attention-pytorch 0.0.3__tar.gz → 0.0.4__tar.gz

{stack_attention_pytorch-0.0.3 → stack_attention_pytorch-0.0.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: stack-attention-pytorch
-Version: 0.0.3
+Version: 0.0.4
 Summary: Implementation of Differentiable Stacks for augmenting Transformers / Attention
 Project-URL: Homepage, https://pypi.org/project/stack-attention-pytorch/
 Project-URL: Repository, https://github.com/lucidrains/stack-attention-pytorch
@@ -63,16 +63,11 @@ tokens = torch.randn(2, 512, 256)
 
 layer = StackTransLayer(256)
 
-out1, state = layer(
-    tokens,
-    stochastic_action = stochastic_action,
-    hard_action = hard_action
-)
+out1, state = layer(tokens)
 
 out2, state = layer(
     tokens,
-    stochastic_action = stochastic_action,
-    hard_action = hard_action
+    stack_states = state
 )
 
 assert out1.shape == out2.shape == tokens.shape

{stack_attention_pytorch-0.0.3 → stack_attention_pytorch-0.0.4}/README.md

@@ -18,16 +18,11 @@ tokens = torch.randn(2, 512, 256)
 
 layer = StackTransLayer(256)
 
-out1, state = layer(
-    tokens,
-    stochastic_action = stochastic_action,
-    hard_action = hard_action
-)
+out1, state = layer(tokens)
 
 out2, state = layer(
     tokens,
-    stochastic_action = stochastic_action,
-    hard_action = hard_action
+    stack_states = state
 )
 
 assert out1.shape == out2.shape == tokens.shape

{stack_attention_pytorch-0.0.3 → stack_attention_pytorch-0.0.4}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "stack-attention-pytorch"
-version = "0.0.3"
+version = "0.0.4"
 description = "Implementation of Differentiable Stacks for augmenting Transformers / Attention"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }

{stack_attention_pytorch-0.0.3 → stack_attention_pytorch-0.0.4}/stack_attention/stack_trans_layer.py

@@ -96,9 +96,9 @@ class StackTransLayer(Module):
         # maybe combining with residual
 
         self.add_residual = add_residual
-        learned_residual_gate &= add_residual
+        self.learned_residual_gate = learned_residual_gate and add_residual
 
-        self.residual_scale = Parameter(tensor(0.)) if learned_residual_gate else None
+        self.residual_scale = Parameter(tensor(0.)) if self.learned_residual_gate else None
 
     def forward(
         self,
@@ -191,7 +191,7 @@ class StackTransLayer(Module):
         read_stack_attn_logits = self.to_read_stack_attn(next_stack_with_null)
         read_stack_attn_logits = rearrange(read_stack_attn_logits, '... 1 -> ...')
 
-        read_stack_attn_logits = read_stack_attn_logits.masked_fill(next_mask_with_null, mask_value(next_stack_with_null))
+        read_stack_attn_logits = read_stack_attn_logits.masked_fill(~next_mask_with_null, mask_value(next_stack_with_null))
 
         read_stack_attn = read_stack_attn_logits.softmax(dim = -1)
 
@@ -199,7 +199,7 @@ class StackTransLayer(Module):
 
         # combine heads
 
-        out = rearrange(stack_inputs, 'b h d -> b (h d)')
+        out = rearrange(read_stack_out, 'b h d -> b (h d)')
 
         out = self.combine(out)
 
@@ -208,6 +208,7 @@ class StackTransLayer(Module):
         # maybe add residual
 
         if self.add_residual:
-            out = out + residual * self.residual_scale.exp()
+            residual_scale = self.residual_scale.exp() if exists(self.residual_scale) else 1.
+            out = out + residual * residual_scale
 
         return out, next_stack_states