torch-einops-utils 0.0.24__py3-none-any.whl → 0.0.25__py3-none-any.whl

torch_einops_utils/__init__.py

@@ -38,7 +38,8 @@ from torch_einops_utils.torch_einops_utils import (
     pad_right_at_dim,
     pad_left_at_dim_to,
     pad_right_at_dim_to,
-    pad_sequence
+    pad_sequence,
+    pad_sequence_and_cat
 )
 
 from torch_einops_utils.torch_einops_utils import (

torch_einops_utils/device.py

@@ -1,8 +1,11 @@
 from itertools import chain
+from functools import wraps
 
 import torch
 from torch.nn import Module
 
+from torch_einops_utils.torch_einops_utils import tree_map_tensor
+
 # helpers
 
 def exists(v):
@@ -18,3 +21,18 @@ def module_device(m: Module):
         return None
 
     return first_param_or_buffer.device
+
+# moving all inputs into a function onto a device
+
+def move_inputs_to_device(device):
+
+    def decorator(fn):
+        @wraps(fn)
+        def inner(*args, **kwargs):
+            args, kwargs = tree_map_tensor(lambda t: t.to(device), (args, kwargs))
+
+            return fn(*args, **kwargs)
+
+        return inner
+
+    return decorator

torch_einops_utils/torch_einops_utils.py

@@ -271,6 +271,17 @@ def pad_sequence(
 
     return output, lens
 
+def pad_sequence_and_cat(
+    tensors,
+    *,
+    dim_cat = 0,
+    **kwargs
+):
+    assert 'return_stacked' not in kwargs
+
+    padded = pad_sequence(tensors, return_stacked = False, **kwargs)
+    return cat(padded, dim = dim_cat)
+
 # tree flatten with inverse
 
 def tree_map_tensor(fn, tree):

{torch_einops_utils-0.0.24.dist-info → torch_einops_utils-0.0.25.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: torch-einops-utils
-Version: 0.0.24
+Version: 0.0.25
 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.25.dist-info/RECORD

@@ -0,0 +1,8 @@
+torch_einops_utils/__init__.py,sha256=RBqs2yPtP51KVi73b-D75S7rIbQJpUOh1IoE3G9j-9E,1008
+torch_einops_utils/device.py,sha256=r4hNXVluxMqk0awCtEMmzbBR5-uJ0KBxPIKGPzVsSwU,794
+torch_einops_utils/save_load.py,sha256=K-i7nmLyXBHdAfBLN3rGQzI3NVf6RRwF_GcrKQnfQsc,2669
+torch_einops_utils/torch_einops_utils.py,sha256=GG1QMs-bH1y3twJ9eu7brmO4fE8tMrmXN_SdfkLD2ZI,6728
+torch_einops_utils-0.0.25.dist-info/METADATA,sha256=FqPrFqbyuWx18fr_TiAsVZj3p_b_NwN6YWaGXs4oIuo,2139
+torch_einops_utils-0.0.25.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+torch_einops_utils-0.0.25.dist-info/licenses/LICENSE,sha256=e6AOF7Z8EFdK3IdcL0x0fLw4cY7Q0d0kNR0o0TmBewM,1066
+torch_einops_utils-0.0.25.dist-info/RECORD,,

torch_einops_utils-0.0.24.dist-info/RECORD

@@ -1,8 +0,0 @@
-torch_einops_utils/__init__.py,sha256=l7S9xJOlbMfwofQD3RI2Z4QxT3H713XSPc1zN_kF4Ww,982
-torch_einops_utils/device.py,sha256=2BFn_RqyNqsThlvjh_Un4bNSkYovUHlNzKF0scq6gaY,366
-torch_einops_utils/save_load.py,sha256=K-i7nmLyXBHdAfBLN3rGQzI3NVf6RRwF_GcrKQnfQsc,2669
-torch_einops_utils/torch_einops_utils.py,sha256=raKFiWpRgGE7Vf7tTqkokU9525gr54lJMpQy_jnUTVI,6498
-torch_einops_utils-0.0.24.dist-info/METADATA,sha256=7WdwsdDrjmI1YuCpyGUkDN0wh1N4X3vyhPbF2mWdK6I,2139
-torch_einops_utils-0.0.24.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-torch_einops_utils-0.0.24.dist-info/licenses/LICENSE,sha256=e6AOF7Z8EFdK3IdcL0x0fLw4cY7Q0d0kNR0o0TmBewM,1066
-torch_einops_utils-0.0.24.dist-info/RECORD,,