torch-einops-utils 0.0.19__py3-none-any.whl → 0.0.21__py3-none-any.whl

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

{torch_einops_utils-0.0.19.dist-info → torch_einops_utils-0.0.21.dist-info}/METADATA

@@ -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.21.dist-info/RECORD

@@ -0,0 +1,7 @@
+torch_einops_utils/__init__.py,sha256=l7S9xJOlbMfwofQD3RI2Z4QxT3H713XSPc1zN_kF4Ww,982
+torch_einops_utils/save_load.py,sha256=K-i7nmLyXBHdAfBLN3rGQzI3NVf6RRwF_GcrKQnfQsc,2669
+torch_einops_utils/torch_einops_utils.py,sha256=exfuRHRBP1YK_3BIpjquFm3fJIX1d9-rNYP1wt9lzoE,6427
+torch_einops_utils-0.0.21.dist-info/METADATA,sha256=1NMFNMuPjmI5xIw4d1r46XtEcotl-OjRsjk8p6lh68k,2139
+torch_einops_utils-0.0.21.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+torch_einops_utils-0.0.21.dist-info/licenses/LICENSE,sha256=e6AOF7Z8EFdK3IdcL0x0fLw4cY7Q0d0kNR0o0TmBewM,1066
+torch_einops_utils-0.0.21.dist-info/RECORD,,

torch_einops_utils-0.0.19.dist-info/RECORD

@@ -1,7 +0,0 @@
-torch_einops_utils/__init__.py,sha256=l7S9xJOlbMfwofQD3RI2Z4QxT3H713XSPc1zN_kF4Ww,982
-torch_einops_utils/save_load.py,sha256=K-i7nmLyXBHdAfBLN3rGQzI3NVf6RRwF_GcrKQnfQsc,2669
-torch_einops_utils/torch_einops_utils.py,sha256=l739ltVspS0u07pMZDqEPG0yYtGnLnwlDQybyri6DnA,6188
-torch_einops_utils-0.0.19.dist-info/METADATA,sha256=N90cpeJGxISPVrjmHLFUcT7DiAt-jXgNZkNmozIxR5U,2139
-torch_einops_utils-0.0.19.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-torch_einops_utils-0.0.19.dist-info/licenses/LICENSE,sha256=e6AOF7Z8EFdK3IdcL0x0fLw4cY7Q0d0kNR0o0TmBewM,1066
-torch_einops_utils-0.0.19.dist-info/RECORD,,