torch-einops-utils 0.0.21__tar.gz → 0.0.23__tar.gz

{torch_einops_utils-0.0.21 → torch_einops_utils-0.0.23}/PKG-INFO

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

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

{torch_einops_utils-0.0.21 → torch_einops_utils-0.0.23}/tests/test_utils.py

@@ -127,6 +127,20 @@ def test_better_pad_sequence():
     mask = lens_to_mask(lens)
     assert torch.allclose(mask.sum(dim = -1), lens)
 
+def test_pad_sequence_uneven_images():
+    images = [
+        torch.randn(3, 16, 17),
+        torch.randn(3, 15, 18),
+        torch.randn(3, 17, 16)
+    ]
+
+    padded_height = pad_sequence(images, dim = -2, return_stacked = False)
+    assert len(padded_height) == 3
+    assert all([t.shape[1] == 17 for t in padded_height])
+
+    stacked = pad_sequence(padded_height, dim = -1, return_stacked = True)
+    assert stacked.shape == (3, 3, 17, 18)
+
 def test_and_masks():
     assert not exists(and_masks([None]))
 

{torch_einops_utils-0.0.21 → torch_einops_utils-0.0.23}/torch_einops_utils/torch_einops_utils.py

@@ -244,6 +244,7 @@ def pad_sequence(
     value = 0.,
     left = False,
     dim_stack = 0,
+    return_stacked = True,
     return_lens = False,
     pad_lens = False # returns padding length instead of sequence lengths
 ):
@@ -258,15 +259,17 @@ def pad_sequence(
     pad_fn = pad_left_at_dim if left else pad_right_at_dim
     padded_tensors = [pad_fn(t, max_len - t_len, dim = dim, value = value) for t, t_len in zip(tensors, lens)]
 
-    stacked = stack(padded_tensors, dim = dim_stack)
+    output = padded_tensors
+    if return_stacked:
+        output = stack(output, dim = dim_stack)
 
     if not return_lens:
-        return stacked
+        return output
 
     if pad_lens:
         lens = max_len - lens
 
-    return stacked, lens
+    return output, lens
 
 # tree flatten with inverse