torch-einops-utils 0.0.15__py3-none-any.whl → 0.0.17__py3-none-any.whl

torch_einops_utils/__init__.py

@@ -1,6 +1,7 @@
 from torch_einops_utils.torch_einops_utils import (
     maybe,
     masked_mean,
+    shape_with_replace,
     slice_at_dim,
     slice_left_at_dim,
     slice_right_at_dim

torch_einops_utils/save_load.py

@@ -0,0 +1,85 @@
+from __future__ import annotations
+from pathlib import Path
+from packaging import version as packaging_version
+
+import pickle
+from functools import wraps
+
+import torch
+from torch.nn import Module
+
+# helpers
+
+def exists(v):
+    return v is not None
+
+def save_load(
+    save_method_name = 'save',
+    load_method_name = 'load',
+    config_instance_var_name = '_config',
+    init_and_load_classmethod_name = 'init_and_load',
+    version: str | None = None
+):
+    def _save_load(klass):
+        assert issubclass(klass, Module), 'save_load should decorate a subclass of torch.nn.Module'
+
+        _orig_init = klass.__init__
+
+        @wraps(_orig_init)
+        def __init__(self, *args, **kwargs):
+            _config = pickle.dumps((args, kwargs))
+
+            setattr(self, config_instance_var_name, _config)
+            _orig_init(self, *args, **kwargs)
+
+        def _save(self, path, overwrite = True):
+            path = Path(path)
+            assert overwrite or not path.exists()
+
+            pkg = dict(
+                model = self.state_dict(),
+                config = getattr(self, config_instance_var_name),
+                version = version,
+            )
+
+            torch.save(pkg, str(path))
+
+        def _load(self, path, strict = True):
+            path = Path(path)
+            assert path.exists()
+
+            pkg = torch.load(str(path), map_location = 'cpu')
+
+            if exists(version) and exists(pkg['version']) and packaging_version.parse(version) != packaging_version.parse(pkg['version']):
+                print(f'loading saved model at version {pkg["version"]}, but current package version is {version}')
+
+            self.load_state_dict(pkg['model'], strict = strict)
+
+        # init and load from
+        # looks for a `config` key in the stored checkpoint, instantiating the model as well as loading the state dict
+
+        @classmethod
+        def _init_and_load_from(cls, path, strict = True):
+            path = Path(path)
+            assert path.exists()
+            pkg = torch.load(str(path), map_location = 'cpu')
+
+            assert 'config' in pkg, 'model configs were not found in this saved checkpoint'
+
+            config = pickle.loads(pkg['config'])
+            args, kwargs = config
+            model = cls(*args, **kwargs)
+
+            _load(model, path, strict = strict)
+            return model
+
+        # set decorated init as well as save, load, and init_and_load
+
+        klass.__init__ = __init__
+        setattr(klass, save_method_name, _save)
+        setattr(klass, load_method_name, _load)
+        setattr(klass, init_and_load_classmethod_name, _init_and_load_from)
+
+        return klass
+
+    return _save_load

torch_einops_utils/torch_einops_utils.py

@@ -61,6 +61,25 @@ def masked_mean(
 
     return num / den.clamp(min = eps)
 
+# shapes
+
+def shape_with_replace(
+    t,
+    replace_dict: dict[int, int] | None = None
+):
+    shape = t.shape
+
+    if not exists(replace_dict):
+        return shape
+
+    shape_list = list(shape)
+
+    for index, value in replace_dict.items():
+        assert index < len(shape_list)
+        shape_list[index] = value
+
+    return torch.Size(shape_list)
+
 # slicing
 
 def slice_at_dim(t, slc, dim = -1):

{torch_einops_utils-0.0.15.dist-info → torch_einops_utils-0.0.17.dist-info}/METADATA

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

@@ -0,0 +1,7 @@
+torch_einops_utils/__init__.py,sha256=F029-BB58UkzkTnXs8odeQWfjkyKnvBBV-7fIg3Drj0,899
+torch_einops_utils/save_load.py,sha256=K-i7nmLyXBHdAfBLN3rGQzI3NVf6RRwF_GcrKQnfQsc,2669
+torch_einops_utils/torch_einops_utils.py,sha256=UpcNBm4XkYF8NCdGGCgD2H1nlikTdtU3JEFWBzMoJjw,5855
+torch_einops_utils-0.0.17.dist-info/METADATA,sha256=OMNDwWMl8oFEH1K1RvL91U-NtjSk5yLSW7XlLIg2iZU,2139
+torch_einops_utils-0.0.17.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+torch_einops_utils-0.0.17.dist-info/licenses/LICENSE,sha256=e6AOF7Z8EFdK3IdcL0x0fLw4cY7Q0d0kNR0o0TmBewM,1066
+torch_einops_utils-0.0.17.dist-info/RECORD,,

torch_einops_utils-0.0.15.dist-info/RECORD

@@ -1,6 +0,0 @@
-torch_einops_utils/__init__.py,sha256=STbYPcW6LEF-ggmMVP_xkwKhyVCN4aHQC-xXGsDN6KA,875
-torch_einops_utils/torch_einops_utils.py,sha256=PZ37JN6Nd7FPYxB_hBgdX6dNZzLpilvtymVtfcF9pFc,5503
-torch_einops_utils-0.0.15.dist-info/METADATA,sha256=XiUbDajKh646ipm8cns-VY0GZVECVMWt3-yoVNhPZpU,2139
-torch_einops_utils-0.0.15.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-torch_einops_utils-0.0.15.dist-info/licenses/LICENSE,sha256=e6AOF7Z8EFdK3IdcL0x0fLw4cY7Q0d0kNR0o0TmBewM,1066
-torch_einops_utils-0.0.15.dist-info/RECORD,,