fusion-bench 0.2.9__py3-none-any.whl

fusion_bench/utils/state_dict_arithmetic.py

@@ -0,0 +1,297 @@
+from collections import OrderedDict
+from numbers import Number
+from typing import Dict, List, Union, cast
+
+import torch
+from torch import Tensor
+
+from .parameters import check_parameters_all_equal
+from .type import StateDictType
+
+
+def to_device(
+    state_dict: StateDictType,
+    device: Union[torch.device, str],
+    copy: bool = False,
+    inplace: bool = False,
+):
+    if inplace:
+        ret_state_dict = state_dict
+    else:
+        ret_state_dict = OrderedDict()
+    for key in state_dict:
+        ret_state_dict[key] = cast(Tensor, state_dict[key]).to(
+            device, non_blocking=True, copy=copy
+        )
+    return ret_state_dict
+
+
+def state_dicts_check_keys(state_dicts: List[StateDictType]):
+    """
+    Checks that the state dictionaries have the same keys.
+
+    Args:
+        state_dicts (List[Dict[str, Tensor]]): A list of dictionaries containing the state of PyTorch models.
+
+    Raises:
+        ValueError: If the state dictionaries have different keys.
+    """
+    # Get the keys of the first state dictionary in the list
+    keys = set(state_dicts[0].keys())
+    # Check that all the state dictionaries have the same keys
+    for state_dict in state_dicts:
+        assert keys == set(state_dict.keys()), "keys of state_dicts are not equal"
+
+
+def num_params_of_state_dict(state_dict: StateDictType):
+    """
+    Returns the number of parameters in a state dict.
+
+    Args:
+        state_dict (Dict[str, Tensor]): The state dict to count the number of parameters in.
+
+    Returns:
+        int: The number of parameters in the state dict.
+    """
+    return sum([state_dict[key].numel() for key in state_dict])
+
+
+def state_dict_flatten(state_dict: Dict[str, Tensor]):
+    """
+    Flattens a state dict.
+
+    Args:
+        state_dict (Dict[str, Tensor]): The state dict to be flattened.
+
+    Returns:
+        Tensor: The flattened state dict.
+    """
+    flattened_state_dict = []
+    for key in state_dict:
+        flattened_state_dict.append(state_dict[key].flatten())
+    return torch.cat(flattened_state_dict)
+
+
+def state_dict_avg(state_dicts: List[StateDictType]):
+    """
+    Returns the average of a list of state dicts.
+
+    Args:
+        state_dicts (List[Dict[str, Tensor]]): The list of state dicts to average.
+
+    Returns:
+        Dict: The average of the state dicts.
+    """
+    assert len(state_dicts) > 0, "The number of state_dicts must be greater than 0"
+    assert all(
+        [len(state_dicts[0]) == len(state_dict) for state_dict in state_dicts]
+    ), "All state_dicts must have the same number of keys"
+
+    num_state_dicts = len(state_dicts)
+    avg_state_dict = OrderedDict()
+    for key in state_dicts[0]:
+        avg_state_dict[key] = torch.zeros_like(state_dicts[0][key])
+        for state_dict in state_dicts:
+            avg_state_dict[key] += state_dict[key]
+        avg_state_dict[key] /= num_state_dicts
+    return avg_state_dict
+
+
+def state_dict_sub(
+    a: StateDictType, b: StateDictType, strict: bool = True, device=None
+):
+    """
+    Returns the difference between two state dicts `a-b`.
+
+    Args:
+        a (StateDictType): The first state dict.
+        b (StateDictType): The second state dict.
+        strict (bool): Whether to check if the keys of the two state dicts are the same.
+
+    Returns:
+        StateDictType: The difference between the two state dicts.
+    """
+    if strict:
+        assert set(a.keys()) == set(b.keys())
+
+    diff = OrderedDict()
+    for k in a:
+        if k in b:
+            diff[k] = a[k] - b[k]
+            if device is not None:
+                diff[k] = diff[k].to(device, non_blocking=True)
+    return diff
+
+
+def state_dict_add(
+    a: StateDictType, b: StateDictType, strict: bool = True, device=None
+):
+    """
+    Returns the sum of two state dicts.
+
+    Args:
+        a (Dict): The first state dict.
+        b (Dict): The second state dict.
+        strict (bool): Whether to check if the keys of the two state dicts are the same.
+
+    Returns:
+        Dict: The sum of the two state dicts.
+    """
+    ans = {}
+    if strict:
+        check_parameters_all_equal([a, b])
+        for key in a:
+            ans[key] = a[key] + b[key]
+    else:
+        for key in a:
+            if key in b:
+                ans[key] = a[key] + b[key]
+    if device is not None:
+        ans = to_device(ans, device)
+    return ans
+
+
+def state_dict_add_scalar(a: StateDictType, scalar: Number):
+    ans = OrderedDict()
+    for key in a:
+        ans[key] = a[key] + scalar
+    return ans
+
+
+def state_dict_mul(state_dict: StateDictType, scalar: float):
+    """
+    Returns the product of a state dict and a scalar.
+
+    Args:
+        state_dict (Dict): The state dict to be multiplied.
+        scalar (float): The scalar to multiply the state dict with.
+
+    Returns:
+        Dict: The product of the state dict and the scalar.
+    """
+    diff = OrderedDict()
+    for k in state_dict:
+        diff[k] = scalar * state_dict[k]
+    return diff
+
+
+def state_dict_div(state_dict: StateDictType, scalar: float):
+    """
+    Returns the division of a state dict by a scalar.
+
+    Args:
+        state_dict (Dict): The state dict to be divided.
+        scalar (float): The scalar to divide the state dict by.
+
+    Returns:
+        Dict: The division of the state dict by the scalar.
+    """
+    diff = OrderedDict()
+    for k in state_dict:
+        diff[k] = state_dict[k] / scalar
+    return diff
+
+
+def state_dict_power(state_dict: Dict[str, Tensor], p: float):
+    """
+    Returns the power of a state dict.
+
+    Args:
+        state_dict (Dict[str, Tensor]): The state dict to be powered.
+        p (float): The power to raise the state dict to.
+
+    Returns:
+        Dict[str, Tensor]: The powered state dict.
+    """
+    powered_state_dict = {}
+    for key in state_dict:
+        powered_state_dict[key] = state_dict[key] ** p
+    return powered_state_dict
+
+
+def state_dict_interpolation(
+    state_dicts: List[Dict[str, Tensor]], scalars: List[float]
+):
+    """
+    Interpolates between a list of state dicts using a list of scalars.
+
+    Args:
+        state_dicts (List[Dict[str, Tensor]]): The list of state dicts to interpolate between.
+        scalars (List[float]): The list of scalars to use for interpolation.
+
+    Returns:
+        Dict: The interpolated state dict.
+    """
+    assert len(state_dicts) == len(
+        scalars
+    ), "The number of state_dicts and scalars must be the same"
+    assert len(state_dicts) > 0, "The number of state_dicts must be greater than 0"
+    assert all(
+        [len(state_dicts[0]) == len(state_dict) for state_dict in state_dicts]
+    ), "All state_dicts must have the same number of keys"
+
+    interpolated_state_dict = {}
+    for key in state_dicts[0]:
+        interpolated_state_dict[key] = torch.zeros_like(state_dicts[0][key])
+        for state_dict, scalar in zip(state_dicts, scalars):
+            interpolated_state_dict[key] += scalar * state_dict[key]
+    return interpolated_state_dict
+
+
+def state_dict_sum(state_dicts: List[StateDictType]):
+    """
+    Returns the sum of a list of state dicts.
+
+    Args:
+        state_dicts (List[Dict[str, Tensor]]): The list of state dicts to sum.
+
+    Returns:
+        Dict: The sum of the state dicts.
+    """
+    assert len(state_dicts) > 0, "The number of state_dicts must be greater than 0"
+    assert all(
+        [len(state_dicts[0]) == len(state_dict) for state_dict in state_dicts]
+    ), "All state_dicts must have the same number of keys"
+
+    sum_state_dict = OrderedDict()
+    for key in state_dicts[0]:
+        sum_state_dict[key] = 0
+        for state_dict in state_dicts:
+            sum_state_dict[key] = sum_state_dict[key] + state_dict[key]
+    return sum_state_dict
+
+
+def state_dict_weighted_sum(
+    state_dicts: List[Dict[str, Tensor]], weights: List[float], device=None
+):
+    """
+    Returns the weighted sum of a list of state dicts.
+
+    Args:
+        state_dicts (List[Dict[str, Tensor]]): The list of state dicts to interpolate between.
+        weights (List[float]): The list of weights to use for the weighted sum.
+
+    Returns:
+        Dict: The weighted sum of the state dicts.
+    """
+    assert len(state_dicts) == len(
+        weights
+    ), "The number of state_dicts and weights must be the same"
+    assert len(state_dicts) > 0, "The number of state_dicts must be greater than 0"
+    assert all(
+        [len(state_dicts[0]) == len(state_dict) for state_dict in state_dicts]
+    ), "All state_dicts must have the same number of keys"
+
+    weighted_sum_state_dict: Dict[str, Tensor] = {}
+    for key in state_dicts[0]:
+        # states dicts can be sparse matrices
+        weighted_sum_state_dict[key] = torch.zeros_like(state_dicts[0][key]).to_dense()
+        for state_dict, weight in zip(state_dicts, weights):
+            weighted_sum_state_dict[key] = torch.add(
+                weighted_sum_state_dict[key], weight * state_dict[key]
+            )
+        if device is not None:
+            weighted_sum_state_dict[key] = weighted_sum_state_dict[key].to(
+                device, non_blocking=True
+            )
+    return weighted_sum_state_dict

fusion_bench/utils/strenum/__init__.py

@@ -0,0 +1,326 @@
+import enum
+
+from ._name_mangler import _NameMangler
+from ._version import get_versions
+
+__version__ = get_versions()["version"]
+__version_info__ = tuple(int(n) for n in __version__.partition("+")[0].split("."))
+del get_versions
+
+_name_mangler = _NameMangler()
+
+# The first argument to the `_generate_next_value_` function of the `enum.Enum`
+# class is documented to be the name of the enum member, not the enum class:
+#
+#     https://docs.python.org/3.6/library/enum.html#using-automatic-values
+#
+# Pylint, though, doesn't know about this so we need to disable it's check for
+# `self` arguments.
+# pylint: disable=no-self-argument
+
+
+class StrEnum(str, enum.Enum):
+    """
+    StrEnum is a Python ``enum.Enum`` that inherits from ``str``. The default
+    ``auto()`` behavior uses the member name as its value.
+
+    Example usage::
+
+        class Example(StrEnum):
+            UPPER_CASE = auto()
+            lower_case = auto()
+            MixedCase = auto()
+
+        assert Example.UPPER_CASE == "UPPER_CASE"
+        assert Example.lower_case == "lower_case"
+        assert Example.MixedCase == "MixedCase"
+    """
+
+    def __new__(cls, value, *args, **kwargs):
+        if not isinstance(value, (str, enum.auto)):
+            raise TypeError(
+                f"Values of StrEnums must be strings: {value!r} is a {type(value)}"
+            )
+        return super().__new__(cls, value, *args, **kwargs)
+
+    def __str__(self):
+        return str(self.value)
+
+    def _generate_next_value_(name, *_):
+        return name
+
+
+class LowercaseStrEnum(StrEnum):
+    """
+    A ``StrEnum`` where ``auto()`` will convert the name to `lowercase` to
+    produce each member's value.
+
+    Example usage::
+
+        class Example(LowercaseStrEnum):
+            UPPER_CASE = auto()
+            lower_case = auto()
+            MixedCase = auto()
+
+        assert Example.UPPER_CASE == "upper_case"
+        assert Example.lower_case == "lower_case"
+        assert Example.MixedCase == "mixedcase"
+
+    .. versionadded:: 0.4.3
+    """
+
+    def _generate_next_value_(name, *_):
+        return name.lower()
+
+
+class UppercaseStrEnum(StrEnum):
+    """
+    A ``StrEnum`` where ``auto()`` will convert the name to `UPPERCASE` to
+    produce each member's value.
+
+    Example usage::
+
+        class Example(UppercaseStrEnum):
+            UPPER_CASE = auto()
+            lower_case = auto()
+            MixedCase = auto()
+
+        assert Example.UPPER_CASE == "UPPER_CASE"
+        assert Example.lower_case == "LOWER_CASE"
+        assert Example.MixedCase == "MIXEDCASE"
+
+    .. versionadded:: 0.4.3
+    """
+
+    def _generate_next_value_(name, *_):
+        return name.upper()
+
+
+class CamelCaseStrEnum(StrEnum):
+    """
+    A ``StrEnum`` where ``auto()`` will convert the name to `camelCase` to
+    produce each member's value.
+
+    Example usage::
+
+        class Example(CamelCaseStrEnum):
+            UPPER_CASE = auto()
+            lower_case = auto()
+            MixedCase = auto()
+
+        assert Example.UPPER_CASE == "upperCase"
+        assert Example.lower_case == "lowerCase"
+        assert Example.MixedCase == "mixedCase"
+
+    .. versionadded:: 0.4.5
+    """
+
+    def _generate_next_value_(name, *_):
+        return _name_mangler.camel(name)
+
+
+class PascalCaseStrEnum(StrEnum):
+    """
+    A ``StrEnum`` where ``auto()`` will convert the name to `PascalCase` to
+    produce each member's value.
+
+    Example usage::
+
+        class Example(PascalCaseStrEnum):
+            UPPER_CASE = auto()
+            lower_case = auto()
+            MixedCase = auto()
+
+        assert Example.UPPER_CASE == "UpperCase"
+        assert Example.lower_case == "LowerCase"
+        assert Example.MixedCase == "MixedCase"
+
+    .. versionadded:: 0.4.5
+    """
+
+    def _generate_next_value_(name, *_):
+        return _name_mangler.pascal(name)
+
+
+class KebabCaseStrEnum(StrEnum):
+    """
+    A ``StrEnum`` where ``auto()`` will convert the name to `kebab-case` to
+    produce each member's value.
+
+    Example usage::
+
+        class Example(KebabCaseStrEnum):
+            UPPER_CASE = auto()
+            lower_case = auto()
+            MixedCase = auto()
+
+        assert Example.UPPER_CASE == "upper-case"
+        assert Example.lower_case == "lower-case"
+        assert Example.MixedCase == "mixed-case"
+
+    .. versionadded:: 0.4.5
+    """
+
+    def _generate_next_value_(name, *_):
+        return _name_mangler.kebab(name)
+
+
+class SnakeCaseStrEnum(StrEnum):
+    """
+    A ``StrEnum`` where ``auto()`` will convert the name to `snake_case` to
+    produce each member's value.
+
+    Example usage::
+
+        class Example(SnakeCaseStrEnum):
+            UPPER_CASE = auto()
+            lower_case = auto()
+            MixedCase = auto()
+
+        assert Example.UPPER_CASE == "upper_case"
+        assert Example.lower_case == "lower_case"
+        assert Example.MixedCase == "mixed_case"
+
+    .. versionadded:: 0.4.5
+    """
+
+    def _generate_next_value_(name, *_):
+        return _name_mangler.snake(name)
+
+
+class MacroCaseStrEnum(StrEnum):
+    """
+    A ``StrEnum`` where ``auto()`` will convert the name to `MACRO_CASE` to
+    produce each member's value.
+
+    Example usage::
+
+        class Example(MacroCaseStrEnum):
+            UPPER_CASE = auto()
+            lower_case = auto()
+            MixedCase = auto()
+
+        assert Example.UPPER_CASE == "UPPER_CASE"
+        assert Example.lower_case == "LOWER_CASE"
+        assert Example.MixedCase == "MIXED_CASE"
+
+    .. versionadded:: 0.4.6
+    """
+
+    def _generate_next_value_(name, *_):
+        return _name_mangler.macro(name)
+
+
+class CamelSnakeCaseStrEnum(StrEnum):
+    """
+    A ``StrEnum`` where ``auto()`` will convert the name to `camel_Snake_Case` to
+    produce each member's value.
+
+    Example usage::
+
+        class Example(CamelSnakeCaseStrEnum):
+            UPPER_CASE = auto()
+            lower_case = auto()
+            MixedCase = auto()
+
+        assert Example.UPPER_CASE == "upper_Case"
+        assert Example.lower_case == "lower_Case"
+        assert Example.MixedCase == "mixed_Case"
+
+    .. versionadded:: 0.4.8
+    """
+
+    def _generate_next_value_(name, *_):
+        return _name_mangler.camel_snake(name)
+
+
+class PascalSnakeCaseStrEnum(StrEnum):
+    """
+    A ``StrEnum`` where ``auto()`` will convert the name to `Pascal_Snake_Case` to
+    produce each member's value.
+
+    Example usage::
+
+        class Example(PascalSnakeCaseStrEnum):
+            UPPER_CASE = auto()
+            lower_case = auto()
+            MixedCase = auto()
+
+        assert Example.UPPER_CASE == "Upper_Case"
+        assert Example.lower_case == "Lower_Case"
+        assert Example.MixedCase == "Mixed_Case"
+
+    .. versionadded:: 0.4.8
+    """
+
+    def _generate_next_value_(name, *_):
+        return _name_mangler.pascal_snake(name)
+
+
+class SpongebobCaseStrEnum(StrEnum):
+    """
+    A ``StrEnum`` where ``auto()`` will convert the name to `SpONGEBob_CAse` to
+    produce each member's value.
+
+    Example usage::
+
+        class Example(SpongebobCaseStrEnum):
+            UPPER_CASE = auto()
+            lower_case = auto()
+            MixedCase = auto()
+
+        assert Example.UPPER_CASE == "uPpER_cAsE"
+        assert Example.lower_case == "lowER_CASe"
+        assert Example.MixedCase == "MixeD_CAse"
+
+    .. versionadded:: 0.4.8
+    """
+
+    def _generate_next_value_(name, *_):
+        return _name_mangler.spongebob(name)
+
+
+class CobolCaseStrEnum(StrEnum):
+    """
+    A ``StrEnum`` where ``auto()`` will convert the name to `COBOL-CASE` to
+    produce each member's value.
+
+    Example usage::
+
+        class Example(CobolCaseStrEnum):
+            UPPER_CASE = auto()
+            lower_case = auto()
+            MixedCase = auto()
+
+        assert Example.UPPER_CASE == "UPPER-CASE"
+        assert Example.lower_case == "LOWER-CASE"
+        assert Example.MixedCase == "MIXED-CASE"
+
+    .. versionadded:: 0.4.8
+    """
+
+    def _generate_next_value_(name, *_):
+        return _name_mangler.cobol(name)
+
+
+class HttpHeaderCaseStrEnum(StrEnum):
+    """
+    A ``StrEnum`` where ``auto()`` will convert the name to `Http-Header-Case` to
+    produce each member's value.
+
+    Example usage::
+
+        class Example(HttpHeaderCaseStrEnum):
+            UPPER_CASE = auto()
+            lower_case = auto()
+            MixedCase = auto()
+
+        assert Example.UPPER_CASE == "Upper-Case"
+        assert Example.lower_case == "Lower-Case"
+        assert Example.MixedCase == "Mixed-Case"
+
+    .. versionadded:: 0.4.8
+    """
+
+    def _generate_next_value_(name, *_):
+        return _name_mangler.http_header(name)