RRAEsTorch 0.1.6__py3-none-any.whl → 0.1.7__py3-none-any.whl

RRAEsTorch/wrappers/__init__.py

@@ -1 +0,0 @@
-from .wrappers import vmap_wrap, norm_wrap

RRAEsTorch/wrappers/wrappers.py

@@ -1,237 +0,0 @@
-import torch
-import dataclasses
-from dataclasses import dataclass
-from torch.func import vmap
-
-@dataclass(frozen=True)
-class NormParams:
-    pass
-
-@dataclass(frozen=True)
-class MeanStdParams(NormParams):
-    mean: float
-    std: float
-
-@dataclass(frozen=True)
-class MinMaxParams(NormParams):
-    min: float
-    max: float
-
-def find_norm_funcs(
-    array=None,
-    norm_typ="None",
-    params=None,
-    ):
-    """ Function that finds norm and inv_norm functions
-
-    The functions can be defined either based on an array (in_train),
-    In this case, the required parameters (e.g., mean, min, etc.) are
-    found using in_train. Otherwise, the parameters can be given
-    explicitly using params_in (check below for some examples).
-
-    Parameters
-    ----------
-    in_train : input array that is to be normalized
-    norm_in: Type of normalization, "minmax" and "meanstd" supported
-    params_in: Normalization parameters if they are to be given manually
-               instead of being computed from in_train.
-               These for example can be:
-               For mean/std normalization, params_out = {"mean": 0.12, "std": 0.23}
-               For min/max normalization, params_out = {"min": -1, "max": 3,}
-    Returns
-    -------
-        A new subclass of base_cls with methods that normalize the input when called.
-    """
-    if norm_typ != "None":
-        assert (params is not None) or (
-            array is not None
-        ), "Either params or in_train must be provided to set norm parameters"
-
-        assert not (
-            params is not None and array is not None
-        ), "Only one of params or in_train must be provided to set norm parameters"
-
-
-    match norm_typ:
-        case "minmax":
-            if params is None:
-                params = MinMaxParams(min=torch.min(array), max=torch.max(array))
-            else:
-                params = params
-            norm_fn = lambda self, x: (x - params.min) / (params.max - params.min)
-            inv_norm_fn = lambda self, x: x * (params.max - params.min) + params.min
-        case "meanstd":
-            if params is None:
-                params = MeanStdParams(mean=torch.mean(array), std=torch.std(array))
-            else:
-                params = params
-            norm_fn = lambda self, x: (x - params.mean) / params.std
-            inv_norm_fn = lambda self, x: x * params.std + params.mean
-        case "None":
-            if params is None:
-                params = NormParams()
-            else:
-                params = params
-            norm_fn = lambda self, x: x
-            inv_norm_fn = lambda self, x: x
-        case _:
-            raise NotImplementedError(f"norm_in specified {norm_typ} is not implemented.")
-
-    return norm_fn, inv_norm_fn, params
-
-def norm_in_wrap(base_cls, array=None, norm_typ="None", params=None, methods_to_wrap=["__call__"]):
-    """ Wrapper that normalizes the input of a function of a subclass of eqx.Module
-
-    The parameters of normalization can be either computed based on an array, or given
-    by the user.
-
-    Parameters
-    ----------
-    base_cls : Base class, subclass of eqx.Module of which functions will
-               be modified.
-    methods_to_wrap: Name of the methods in base_cls to be wrapped
-    array : The array from which norm parameters are to be found
-    norm_typ: Type of normalization, "meanstd" and "minmax" are supported
-    params: Parameters of normalization if these are to be given manually
-                These for example can be:
-                For mean/std normalization, params_out = {"mean": 0.12, "std": 0.23}
-                For min/max normalization, params_out = {"min": -1, "max": 3,}
-    Returns
-    -------
-        A new subclass of base_cls with methods that normalize the input when called.
-    """
-    norm_in, inv_norm_in, params_in = find_norm_funcs(array, norm_typ, params)
-    def norm_in_decorator(fn):
-        def wrapped(self, x, *args, **kwargs):
-            result = fn(self, norm_in(self, x), *args, **kwargs)
-            return result
-        return wrapped, {"norm_in": (callable, norm_in), "inv_norm_in": (callable, inv_norm_in), "params_in": (dict, params_in)}
-    return make_wrapped(base_cls, norm_in_decorator, methods_to_wrap)
-
-def inv_norm_out_wrap(base_cls, array=None, norm_typ="None", params=None, methods_to_wrap=["__call__"]):
-    """ Wrapper that de-normalizes the output of a function of a subclass of eqx.Module
-
-    The parameters of normalization can be either computed based on an array, or given
-    by the user.
-
-    Parameters
-    ----------
-    base_cls : Base class, subclass of eqx.Module of which functions will
-               be modified.
-    methods_to_wrap: Name of the methods in base_cls to be wrapped
-    array : The array from which norm parameters are to be found
-    norm_typ: Type of normalization, "meanstd" and "minmax" are supported
-    params: Parameters of normalization if these are to be given manually
-                These for example can be:
-                For mean/std normalization, params_out = {"mean": 0.12, "std": 0.23}
-                For min/max normalization, params_out = {"min": -1, "max": 3,}
-    Returns
-    -------
-        A new subclass of base_cls with de-normalized methods.
-        The given methods accept an additional keyword argument "keep_normalized"
-        If this argument is passed as True, functions retrieve original behavior
-    """
-    norm_out, inv_norm_out, params_out = find_norm_funcs(array, norm_typ, params)
-
-    def norm_out_decorator(fn):
-        def wrapped(self, x, *args, keep_normalized=False, **kwargs):
-            if keep_normalized:
-                result = fn(self, x, *args, **kwargs)
-            else:
-                result = inv_norm_out(self, fn(self, x, *args, **kwargs))
-            return result
-        return wrapped, {"norm_out": (callable, norm_out), "inv_norm_out": (callable, inv_norm_out), "params_out": (dict, params_out)}
-    return make_wrapped(base_cls, norm_out_decorator, methods_to_wrap)
-
-def norm_wrap(base_cls, array_in=None, norm_typ_in="None", params_in=None, array_out=None, norm_typ_out="None", params_out=None, methods_to_wrap_in=["__call__"], methods_to_wrap_out=["__call__"]):
-    """ Wrapper that normalizes functions of a subclass of eqx.Module
-
-    Parameters
-    ----------
-    base_cls : Base class, subclass of eqx.Module of which functions will
-               be modified.
-    methods_to_wrap: Name of the methods in base_cls to be wrapped
-    ... : Other parameters are explained in norm_in_wrap and norm_out_wrap
-
-    Returns
-    -------
-        A new subclass of base_cls with normalized methods.
-    """
-    after_in = norm_in_wrap(base_cls, array_in, norm_typ_in, params_in, methods_to_wrap_in)
-    after_out = inv_norm_out_wrap(after_in, array_out, norm_typ_out, params_out, methods_to_wrap_out)
-    return after_out
-
-def vmap_wrap(base_cls, map_axis, count=1, methods_to_wrap=["__call__"]):
-    """ Wrapper that vectorizes functions of a subclass of eqx.Module
-
-    Parameters
-    ----------
-    base_cls : Base class, subclass of eqx.Module of which functions will
-               be modified.
-    map_axis : Axis along which to vectorize the functions
-    methods_to_wrap: Name of the methods in base_cls to be wrapped
-    count: How many times to vectorize the functions
-
-    Returns
-    -------
-        A new subclass of base_cls with vectorized methods.
-        The given methods accept an additional keyword argument "no_map"
-        If this argument is passed as True, functions retrieve original behavior
-    """
-    def vmap_decorator(fn):
-        def wrapped(self, x, *args, no_map=False, **kwargs):
-            if (map_axis is None) or no_map:
-                return fn(self, x, *args, **kwargs)
-            f = lambda x: fn(self, x, *args, **kwargs)
-            for _ in range(count):
-                f = vmap(f, in_dims=(map_axis,), out_dims=map_axis)
-            out = f(x)
-            return out
-        return wrapped, {}
-    return make_wrapped(base_cls, vmap_decorator, methods_to_wrap)
-
-
-
-def make_wrapped(base_cls, decorator, methods_to_wrap=["__call__"]):
-    """
-    Create a subclass of base_cls with specified methods wrapped by decorator.
-
-    Parameters
-    ----------
-        base_cls: Original class to wrap.
-        methods_to_wrap: List of method names (strings) to decorate.
-        decorator: The wanted modification to the methods given above
-
-    Returns
-    -------
-        A new subclass of base_cls with decorated methods.
-    """
-    attrs = {}
-    annotations = {}
-    seen_fields = set()
-
-    for method_name in methods_to_wrap:
-        if not hasattr(base_cls, method_name):
-            raise AttributeError(f"Method {method_name} not found in {base_cls}")
-        
-        original = getattr(base_cls, method_name)
-        wrapped_method, extra_fields = decorator(original)
-        attrs[method_name] = wrapped_method
-
-        # Sort field names to avoid ordering issues
-        for field_name in extra_fields:  
-            if field_name in seen_fields:
-                continue
-
-            field_type, default_value = extra_fields[field_name]
-            annotations[field_name] = field_type
-
-            attrs[field_name] = dataclasses.field(default=default_value)
-
-            seen_fields.add(field_name)
-
-    if annotations:
-        attrs["__annotations__"] = annotations
-
-    Wrapped = type(base_cls.__name__, (base_cls,), attrs)
-    return Wrapped

rraestorch-0.1.6.dist-info/RECORD

@@ -1,26 +0,0 @@
-RRAEsTorch/__init__.py,sha256=f234R6usRCqIgmBmiXyZNIHa7VrDe5E-KZO0Y6Ek5AQ,33
-RRAEsTorch/config.py,sha256=bQPwc_2KTvhglH_WIRSb5_6CpUQQj9AGpfqBp8_kuys,2931
-RRAEsTorch/AE_base/AE_base.py,sha256=Eeo_I7p5P-357rnOmCuFxosJgmBg4KPyMA8n70sTV7U,3368
-RRAEsTorch/AE_base/__init__.py,sha256=95YfMgEWzIFAkm--Ci-a9YPSGfCs2PDAK2sbfScT7oo,24
-RRAEsTorch/AE_classes/AE_classes.py,sha256=oDpDzQasPbtK2L9vDLiG4VQdKH02VRCagOYT1-FAldo,18063
-RRAEsTorch/AE_classes/__init__.py,sha256=inM2_YPJG8T-lwx-CUg-zL2EMltmROQAlNZeZmnvVGA,27
-RRAEsTorch/tests/test_AE_classes_CNN.py,sha256=bEE9JnTo84t9w0a4kw1W74L51eLGjBB8trrlAG938RE,3182
-RRAEsTorch/tests/test_AE_classes_MLP.py,sha256=Cr1_uP7lag6RPQC1UhN2O7RFW4BEx1cd0Z-Y6VgrWRg,2718
-RRAEsTorch/tests/test_fitting_CNN.py,sha256=8i6oUZFS-DpaTh1VsRsd1rGG_Me7R3Kf1vO4lPSxYns,2797
-RRAEsTorch/tests/test_fitting_MLP.py,sha256=6Ggy5iAJmqMEJPDoon7fYkwFS8LWmosTOFBQuAQzszc,3353
-RRAEsTorch/tests/test_mains.py,sha256=ivTXP7NypSlgmB9VR5g0yq5VEuPZJGOibDqBMjOxHow,1021
-RRAEsTorch/tests/test_save.py,sha256=_zYcDZtz_HWzHHWT6uJVq7ynsFmlH5v6nDE-loZQ4zo,1997
-RRAEsTorch/tests/test_stable_SVD.py,sha256=OimHPqw4f22qndyRzwJfNvTzzjP2CM-yHtfXCqkMBuA,1230
-RRAEsTorch/tests/test_wrappers.py,sha256=Ike4IfMUx2Qic3f3_cBikgFPEU1WW5TuH1jT_r2NgvY,2215
-RRAEsTorch/trackers/__init__.py,sha256=3c9qcUMZiUfVr93rxFp6l11lIDthyK3PCY_-P-sNX3I,25
-RRAEsTorch/trackers/trackers.py,sha256=Pn1ejMxMjAtvgDazFFwa3qiZhogG5GtXj4UIIFiBpuY,9127
-RRAEsTorch/training_classes/__init__.py,sha256=K_Id4yhw640jp2JN15-0E4wJi4sPadi1fFRgovMV3kw,101
-RRAEsTorch/training_classes/training_classes.py,sha256=HU8Ksz1-2WwOMuwyGiWdkQ_vrrgBEwaeQT4avs4jd2E,37870
-RRAEsTorch/utilities/__init__.py,sha256=NtlizCcRW4qcsULXxWfjPk265rLJst0-GqWLRah2yDY,26
-RRAEsTorch/utilities/utilities.py,sha256=JfLkAPEC8fzwgM32LEcXVe0tA4C7UBgsrkuh6noUA_4,53372
-RRAEsTorch/wrappers/__init__.py,sha256=txiLh4ylnuvPlapagz7DiAslmjllOzTqwCDL2dFr6dM,44
-RRAEsTorch/wrappers/wrappers.py,sha256=9Rmq2RS_EkZvsg96SKrt1HFIP35sF0xyPI0goV0ujOs,9659
-rraestorch-0.1.6.dist-info/METADATA,sha256=BfnB-vhx0m-d79hCd8UgZT0c8GPHFRwVn8x-M8k_h6E,3028
-rraestorch-0.1.6.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-rraestorch-0.1.6.dist-info/licenses/LICENSE,sha256=QQKXj7kEw2yUJxPDzHUnEPRO7YweIU6AAlwPLnhtinA,1090
-rraestorch-0.1.6.dist-info/RECORD,,