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

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

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

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

torch_einops_utils-0.0.25/tests/test_device.py

@@ -0,0 +1,17 @@
+import torch
+from torch import nn
+
+def test_module_device():
+    from torch_einops_utils.device import module_device
+
+    assert module_device(nn.Linear(3, 3)) == torch.device('cpu')
+    assert module_device(nn.Identity()) is None
+
+def test_move_input_to_device():
+    from torch_einops_utils.device import move_inputs_to_device
+
+    def fn(t):
+        return t
+
+    decorated = move_inputs_to_device(torch.device('cpu'))(fn)
+    decorated(torch.randn(3))

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

@@ -17,6 +17,7 @@ from torch_einops_utils.torch_einops_utils import (
     pad_left_at_dim_to,
     pad_right_at_dim_to,
     pad_sequence,
+    pad_sequence_and_cat,
     lens_to_mask,
     and_masks,
     or_masks,
@@ -138,8 +139,8 @@ def test_pad_sequence_uneven_images():
     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)
+    stacked = pad_sequence_and_cat(padded_height, dim_cat = 0)
+    assert stacked.shape == (9, 17, 18)
 
 def test_and_masks():
     assert not exists(and_masks([None]))

{torch_einops_utils-0.0.23 → torch_einops_utils-0.0.25}/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-0.0.25/torch_einops_utils/device.py

@@ -0,0 +1,38 @@
+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):
+    return v is not None
+
+# infer the device for a module
+
+def module_device(m: Module):
+
+    first_param_or_buffer = next(chain(m.parameters(), m.buffers()), None)
+
+    if not exists(first_param_or_buffer):
+        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-0.0.23 → torch_einops_utils-0.0.25}/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):