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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: torch-einops-utils
-Version: 0.1.0
+Version: 0.1.2
 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.1.0 → torch_einops_utils-0.1.2}/pyproject.toml

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

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

@@ -310,3 +310,31 @@ def test_shift():
 def test_reverse_cumsum():
     t = tensor([1, 2, 3])
     assert reverse_cumsum(t).tolist() == [6, 5, 3]
+
+def test_pad_right_ndim_to_and_expand_as():
+    from torch_einops_utils.torch_einops_utils import pad_right_ndim_to_and_expand_as
+
+    target = torch.randn(2, 8, 64)
+    source = torch.randint(0, 8, (2, 4))
+    assert pad_right_ndim_to_and_expand_as(source, target).shape == (2, 4, 64)
+
+    dest = torch.zeros(2, 8, 64)
+    source = torch.arange(4).unsqueeze(0).expand(2, -1)
+
+    scattered = dest.scatter(1, pad_right_ndim_to_and_expand_as(source, dest), torch.ones(2, 4, 64))
+
+    assert (scattered[:, :4] == 1.).all()
+    assert (scattered[:, 4:] == 0.).all()
+
+def test_repeat_interleave_to_match():
+    from torch_einops_utils.torch_einops_utils import repeat_interleave_to_match
+    time_lens = torch.tensor([2, 3])
+
+    out = repeat_interleave_to_match(time_lens, torch.randn(4, 512))
+    assert out.tolist() == [2, 2, 3, 3]
+
+    out2 = repeat_interleave_to_match(time_lens, torch.randn(6, 128))
+    assert out2.tolist() == [2, 2, 2, 3, 3, 3]
+
+    out3 = repeat_interleave_to_match(time_lens, 6)
+    assert out3.tolist() == [2, 2, 2, 3, 3, 3]

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

@@ -57,6 +57,11 @@ from torch_einops_utils.torch_einops_utils import (
     pack_with_inverse
 )
 
+from torch_einops_utils.torch_einops_utils import (
+    pad_right_ndim_to_and_expand_as,
+    repeat_interleave_to_match
+)
+
 from torch_einops_utils.nn import (
     Sequential,
     Lambda,

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

@@ -17,6 +17,9 @@ def exists(v):
 def default(v, d):
     return v if exists(v) else d
 
+def divisible_by(num, den):
+    return (num % den) == 0
+
 def identity(t, *args, **kwargs):
     return t
 
@@ -349,3 +352,25 @@ def pack_with_inverse(t, pattern):
         return out
 
     return packed, inverse
+
+# gather and scatter
+
+def pad_right_ndim_to_and_expand_as(source, target):
+    source_ndim = source.ndim
+    source = pad_right_ndim_to(source, target.ndim)
+    shape = [*source.shape[:source_ndim], *target.shape[source_ndim:]]
+    return source.expand(shape)
+
+# repeat
+
+def repeat_interleave_to_match(t, target, dim = 0, target_dim = None):
+    target_dim = default(target_dim, dim)
+
+    len_t = t.shape[dim]
+    len_target = target if isinstance(target, int) else target.shape[target_dim]
+
+    if len_t == len_target:
+        return t
+
+    assert divisible_by(len_target, len_t)
+    return t.repeat_interleave(len_target // len_t, dim = dim)