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

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

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

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

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

@@ -210,6 +210,7 @@ def test_safe_functions():
     assert safe_stack([None]) is None
     assert (safe_stack([t1]) == t1).all()
     assert (safe_stack([t1, None]) == t1).all()
+    assert safe_stack([t1]).shape == (1, 2, 3)
     assert safe_stack([t1, t2]).shape == (2, 2, 3)
 
     assert safe_cat([]) is None

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

@@ -23,6 +23,9 @@ def identity(t, *args, **kwargs):
 def first(arr):
     return arr[0]
 
+def compact(arr):
+    return [*filter(exists, arr)]
+
 def maybe(fn):
 
     if not exists(fn):
@@ -40,7 +43,7 @@ def maybe(fn):
 def safe(fn):
     @wraps(fn)
     def inner(tensors, *args, **kwargs):
-        tensors = [*filter(exists, tensors)]
+        tensors = compact(tensors)
 
         if len(tensors) == 0:
             return None
@@ -160,8 +163,12 @@ def align_dims_left(
 
 # cat and stack
 
-@safe
 def safe_stack(tensors, dim = 0):
+    tensors = compact(tensors)
+
+    if len(tensors) == 0:
+        return None
+
     return stack(tensors, dim = dim)
 
 @safe
@@ -237,7 +244,8 @@ def pad_sequence(
     value = 0.,
     left = False,
     dim_stack = 0,
-    return_lens = False
+    return_lens = False,
+    pad_lens = False # returns padding length instead of sequence lengths
 ):
     if len(tensors) == 0:
         return None
@@ -255,6 +263,9 @@ def pad_sequence(
     if not return_lens:
         return stacked
 
+    if pad_lens:
+        lens = max_len - lens
+
     return stacked, lens
 
 # tree flatten with inverse