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

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

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

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

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

@@ -325,3 +325,16 @@ def test_pad_right_ndim_to_and_expand_as():
 
     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.1 → torch_einops_utils-0.1.2}/torch_einops_utils/__init__.py

@@ -58,7 +58,8 @@ from torch_einops_utils.torch_einops_utils import (
 )
 
 from torch_einops_utils.torch_einops_utils import (
-    pad_right_ndim_to_and_expand_as
+    pad_right_ndim_to_and_expand_as,
+    repeat_interleave_to_match
 )
 
 from torch_einops_utils.nn import (

{torch_einops_utils-0.1.1 → 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
 
@@ -357,3 +360,17 @@ def pad_right_ndim_to_and_expand_as(source, target):
     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)