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

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

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

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

torch_einops_utils-0.0.24/tests/test_device.py

@@ -0,0 +1,8 @@
+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

torch_einops_utils-0.0.24/torch_einops_utils/device.py

@@ -0,0 +1,20 @@
+from itertools import chain
+
+import torch
+from torch.nn import Module
+
+# 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