torch-einops-utils 0.0.34__tar.gz → 0.1.0__tar.gz

{torch_einops_utils-0.0.34 → torch_einops_utils-0.1.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: torch-einops-utils
-Version: 0.0.34
+Version: 0.1.0
 Summary: Personal utility functions
 Project-URL: Homepage, https://pypi.org/project/torch-einops-utils/
 Project-URL: Repository, https://github.com/lucidrains/torch-einops-utils

{torch_einops_utils-0.0.34 → torch_einops_utils-0.1.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "torch-einops-utils"
-version = "0.0.34"
+version = "0.1.0"
 description = "Personal utility functions"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }

{torch_einops_utils-0.0.34 → torch_einops_utils-0.1.0}/tests/test_nn.py

@@ -10,7 +10,7 @@ def test_sequential():
         nn.ReLU()
     )
     assert len(seq) == 2
-    
+
     # Test forward pass
     x = torch.randn(2, 10)
     out = seq(x)

{torch_einops_utils-0.0.34 → torch_einops_utils-0.1.0}/tests/test_save_load.py

@@ -19,22 +19,22 @@ class SimpleNet(nn.Module):
 def test_save_load():
     model = SimpleNet(10, 20)
     path = Path('test_model.pt')
-    
+
     # Save the model
     model.save(str(path))
-    
+
     # Create another model with different weights
     model2 = SimpleNet(10, 20)
-    
+
     # Ensure weights are different initially
     assert not torch.allclose(model.net.weight, model2.net.weight)
-    
+
     # Load back
     model2.load(str(path))
-    
+
     # Validate weights are the same
     assert torch.allclose(model.net.weight, model2.net.weight)
-    
+
     # Cleanup
     if path.exists():
         os.remove(path)
@@ -43,18 +43,18 @@ def test_init_and_load():
     dim, hidden_dim = 16, 32
     model = SimpleNet(dim, hidden_dim)
     path = Path('test_model_init.pt')
-    
+
     # Save the model
     model.save(str(path))
-    
+
     # Initialize and load from file
     model2 = SimpleNet.init_and_load(str(path))
-    
+
     # Validate attributes and weights
     assert model2.dim == dim
     assert model2.hidden_dim == hidden_dim
     assert torch.allclose(model.net.weight, model2.net.weight)
-    
+
     # Cleanup
     if path.exists():
         os.remove(path)
@@ -97,32 +97,32 @@ def test_sophisticated_nested_save_load():
     c1 = Child(name = "c1")
     c2 = Child(name = "c2")
     c_nest = Child(grandchild = gc, name = "c_nest")
-    
+
     p1 = Parent(child1 = c1, child2 = c2)
     p2 = Parent(child1 = c_nest)
-    
+
     gp = GrandParent(p1 = p1, p2 = p2)
-    
+
     path = Path('sophisticated_test.pt')
-    
+
     # Save
     gp.save(str(path))
-    
+
     # Load
     gp2 = GrandParent.init_and_load(str(path))
-    
+
     # Verify structure
     assert gp2.p1.child1.name == "c1"
     assert gp2.p1.child2.name == "c2"
     assert gp2.p2.child1.name == "c_nest"
     assert gp2.p2.child1.grandchild.dim == 8
-    
+
     # Verify weight parity
     assert torch.allclose(gp.param, gp2.param)
     assert torch.allclose(gp.p1.param, gp2.p1.param)
     assert torch.allclose(gp.p1.child1.param, gp2.p1.child1.param)
     assert torch.allclose(gp.p2.child1.grandchild.param, gp2.p2.child1.grandchild.param)
- 
+
     if path.exists():
         os.remove(path)
 

{torch_einops_utils-0.0.34 → torch_einops_utils-0.1.0}/tests/test_utils.py

@@ -30,7 +30,12 @@ from torch_einops_utils.torch_einops_utils import (
     slice_left_at_dim,
     slice_right_at_dim,
     safe_stack,
-    safe_cat
+    safe_cat,
+    mask_after,
+    mask_before,
+    shift_right,
+    shift_left,
+    reverse_cumsum
 )
 
 def test_exist():
@@ -250,3 +255,58 @@ def test_safe_functions():
     assert (safe_cat([t1]) == t1).all()
     assert (safe_cat([t1, None]) == t1).all()
     assert safe_cat([t1, t2]).shape == (4, 3)
+
+def test_mask_after_before():
+    t = tensor([[1, 2, 3, 4, 5], [1, 3, 2, 3, 5]])
+
+    assert mask_after(t, 3).tolist() == [
+        [True, True, True, False, False],
+        [True, True, False, False, False]
+    ]
+
+    assert mask_after(t, 3, inclusive = False).tolist() == [
+        [True, True, False, False, False],
+        [True, False, False, False, False]
+    ]
+
+    assert mask_before(t, 3).tolist() == [
+        [False, False, True, True, True],
+        [False, False, False, True, True]
+    ]
+
+    assert mask_before(t, 3, inclusive = False).tolist() == [
+        [False, False, False, True, True],
+        [False, False, False, False, True]
+    ]
+
+    assert mask_after(t.T, 3, dim = 0).tolist() == mask_after(t, 3).T.tolist()
+    assert mask_before(t.T, 3, dim = 0).tolist() == mask_before(t, 3).T.tolist()
+
+def test_eos_id_masking():
+    seq = tensor([
+        [1, 4, 5, 2, 0, 0],
+        [1, 6, 2, 0, 0, 0],
+        [1, 7, 8, 9, 2, 0]
+    ])
+
+    assert mask_after(seq, 2).tolist() == [
+        [True, True, True, True, False, False],
+        [True, True, True, False, False, False],
+        [True, True, True, True, True, False]
+    ]
+
+    assert mask_after(seq, 2, inclusive = False).tolist() == [
+        [True, True, True, False, False, False],
+        [True, True, False, False, False, False],
+        [True, True, True, True, False, False]
+    ]
+
+def test_shift():
+    t = tensor([1, 2, 3])
+    assert shift_right(t).tolist() == [0, 1, 2]
+    assert shift_left(t).tolist() == [2, 3, 0]
+    assert shift_right(t, pad_value = -1).tolist() == [-1, 1, 2]
+
+def test_reverse_cumsum():
+    t = tensor([1, 2, 3])
+    assert reverse_cumsum(t).tolist() == [6, 5, 3]

{torch_einops_utils-0.0.34 → torch_einops_utils-0.1.0}/torch_einops_utils/__init__.py

@@ -20,7 +20,9 @@ from torch_einops_utils.torch_einops_utils import (
     lens_to_mask,
     reduce_masks,
     and_masks,
-    or_masks
+    or_masks,
+    mask_after,
+    mask_before
 )
 
 from torch_einops_utils.torch_einops_utils import (
@@ -29,7 +31,8 @@ from torch_einops_utils.torch_einops_utils import (
 )
 
 from torch_einops_utils.torch_einops_utils import (
-    exclusive_cumsum
+    exclusive_cumsum,
+    reverse_cumsum
 )
 
 from torch_einops_utils.torch_einops_utils import (
@@ -39,7 +42,10 @@ from torch_einops_utils.torch_einops_utils import (
     pad_left_at_dim_to,
     pad_right_at_dim_to,
     pad_sequence,
-    pad_sequence_and_cat
+    pad_sequence_and_cat,
+    shift,
+    shift_right,
+    shift_left
 )
 
 from torch_einops_utils.torch_einops_utils import (

{torch_einops_utils-0.0.34 → torch_einops_utils-0.1.0}/torch_einops_utils/torch_einops_utils.py

@@ -87,6 +87,9 @@ def masked_mean(
 def exclusive_cumsum(t, dim = -1):
     return t.cumsum(dim = dim) - t
 
+def reverse_cumsum(t, dim = -1):
+    return t.sum(dim = dim, keepdim = True) - t.cumsum(dim = dim) + t
+
 # shapes
 
 def shape_with_replace(
@@ -210,6 +213,18 @@ def and_masks(masks):
 def or_masks(masks):
     return reduce_masks(masks, torch.logical_or)
 
+def mask_after(t, value, dim = -1, inclusive = True):
+    mask = t == value
+    if inclusive:
+        mask = shift_right(mask, amount = 1, dim = dim, pad_value = False)
+    return mask.float().cumsum(dim = dim) == 0.
+
+def mask_before(t, value, dim = -1, inclusive = True):
+    mask = t == value
+    if inclusive:
+        mask = shift_left(mask, amount = 1, dim = dim, pad_value = False)
+    return reverse_cumsum(mask.float(), dim = dim) == 0.
+
 # padding
 
 def pad_at_dim(
@@ -243,6 +258,17 @@ def pad_right_at_dim_to(t, length: int, dim = -1, **kwargs):
 
     return pad_right_at_dim(t, length - curr_len, dim = dim, **kwargs)
 
+# shifting
+
+def shift(t, amount = 1, dim = -1, pad_value = 0.):
+    return pad_at_dim(t, (amount, -amount), dim = dim, value = pad_value)
+
+def shift_right(t, amount = 1, dim = -1, pad_value = 0.):
+    return shift(t, amount, dim = dim, pad_value = pad_value)
+
+def shift_left(t, amount = 1, dim = -1, pad_value = 0.):
+    return shift(t, -amount, dim = dim, pad_value = pad_value)
+
 # better pad sequence
 
 def pad_sequence(