Prompt:
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import fnmatch
import logging
from typing import Any, Callable, Dict, List, Optional, Set, Type, Union

import hydra
import torch
import torch.nn as nn
from iopath.common.file_io import g_pathmgr
from omegaconf import OmegaConf

from .model_wrappers import MIMOHeadWrapper

def _unix_pattern_to_parameter_names(
    constraints: List[str], all_parameter_names: Set[str]
) -> Union[None, Set[str]]:

    parameter_names = []
    for param_name in constraints:
        matching_parameters = set(fnmatch.filter(all_parameter_names, param_name))
        assert (
            len(matching_parameters) > 0
        ), f"param_names {param_name} don't match any param in the given names."
        parameter_names.append(matching_parameters)
    return set.union(*parameter_names)

class CkptIncludeKernel:
    """
    Includes only the keys from the given model state_dict that match the key_pattern.
    Rest of the keys are removed from the given state_dict.

    Args:
        key_pattern: Patterns used to select the keys in the state_dict
            that are eligible for this kernel.
    """

    def __init__(self, key_pattern: List[str]):
        self.key_pattern = key_pattern

    def __call__(self, state_dict: Dict):
        """
        Args:
            state_dict: A dictionary representing the given checkpoint's state dict.
        """

        include_keys = _unix_pattern_to_parameter_names(
            self.key_pattern, state_dict.keys()
        )

        new_state_dict = {}
        for key in include_keys:
            new_state_dict[key] = state_dict[key]

        return new_state_dict

class CkptExcludeKernel:
    """
    Removes the keys from the given model state_dict that match the key_pattern.

    Args:
        key_pattern: Patterns used to select the keys in the state_dict
            that are eligible for this kernel.
    """

    def __init__(self, key_pattern: List[str]):
        self.key_pattern = key_pattern

    def __call__(self, state_dict: Dict):
        """
        Args:
            state_dict: A dictionary representing the given checkpoint's state dict.
        """

        exclude_keys = _unix_pattern_to_parameter_names(
            self.key_pattern, state_dict.keys()
        )
        include_keys = set(state_dict.keys()) - exclude_keys

        new_state_dict = {}
        for key in include_keys:
            new_state_dict[key] = state_dict[key]

        return new_state_dict

class CkptPrependKernel:
    """
    Prepends the given pattern to all the keys in the checkpoint state dict after
    selecting them with key_pattern.

    For instance, if prepend_pattern  = "some_prepend." and
    key_pattern = ["model.head"], this kernel would prepend "some_prepend." to
    "model.key", thus renaming the key "model.head" to "some_prepend.model.head".

    Args:
        prepend_pattern: The pattern to prepend the keys in the state_dict with.
        key_pattern: Patterns used to select the keys in the state_dict
            that are eligible for this kernel.
    """

    def __init__(self, prepend_pattern: str, key_pattern: List[str]):
        self.prepend_pattern = prepend_pattern
        self.key_pattern = key_pattern

    def __call__(self, state_dict: Dict):
        """
        Args:
            state_dict: A dictionary representing the given checkpoint's state dict.
        """

        all_keys = set(state_dict.keys())

        include_keys = set(state_dict.keys())
        if self.key_pattern is not None:
            include_keys = _unix_pattern_to_parameter_names(
                self.key_pattern, state_dict.keys()
            )

        excluded_keys = all_keys - include_keys

        # Add excluded keys from re-mapping
        new_state_dict = {}
        for k in excluded_keys:
            new_state_dict[k] = state_dict[k]

        # Add keys from remapping
        for key in include_keys:
            new_state_dict[self.prepend_pattern + key] = state_dict[key]

        return new_state_dict

class CkptRenameWithCopyKernel:
    """
    Renames and also optionally creates copyies of the key-value pairs in the checkpoint
    state dict. Before doing so, selects the keys to which to apply this kernel by
    using key_pattern.

    For instance, if source_pattern  = "model.head" and
    target_patterns = ["model.head_1", "model.head_2"], this kernel would
    rename the key "model.head" to "model.head_1" and will also create a copy of the
    "model.head" and assign it a new name "model.head_2".

    Args:
        source_pattern: The pattern that needs to be renamed in the current
            checkpoint state_dict.
        target_patterns: A list of patterns to which the source_pattern is to be
            renamed to it. If the list has more than one element, it creates multiple
            copies of the source_pattern value and assigns then the names given in
            target_pattern.
        key_pattern: Patterns used to select the keys in the state_dict
            that are eligible for this kernel.
    """

    def __init__(
        self,
        source_pattern: str,
        target_patterns: List[str],
        key_pattern: Optional[List[str]] = None,
    ):

        self.source_pattern = source_pattern
        self.target_patterns = target_patterns
        self.key_pattern = key_pattern


Prompt:
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import fnmatch
import logging
from typing import Any, Callable, Dict, List, Optional, Set, Type, Union

import hydra
import torch
import torch.nn as nn
from iopath.common.file_io import g_pathmgr
from omegaconf import OmegaConf

from .model_wrappers import MIMOHeadWrapper

def _unix_pattern_to_parameter_names(
    constraints: List[str], all_parameter_names: Set[str]
) -> Union[None, Set[str]]:

    parameter_names = []
    for param_name in constraints:
        matching_parameters = set(fnmatch.filter(all_parameter_names, param_name))
        assert (
            len(matching_parameters) > 0
        ), f"param_names {param_name} don't match any param in the given names."
        parameter_names.append(matching_parameters)
    return set.union(*parameter_names)

class CkptIncludeKernel:
    """
    Includes only the keys from the given model state_dict that match the key_pattern.
    Rest of the keys are removed from the given state_dict.

    Args:
        key_pattern: Patterns used to select the keys in the state_dict
            that are eligible for this kernel.
    """

    def __init__(self, key_pattern: List[str]):
        self.key_pattern = key_pattern

    def __call__(self, state_dict: Dict):
        """
        Args:
            state_dict: A dictionary representing the given checkpoint's state dict.
        """

        include_keys = _unix_pattern_to_parameter_names(
            self.key_pattern, state_dict.keys()
        )

        new_state_dict = {}
        for key in include_keys:
            new_state_dict[key] = state_dict[key]

        return new_state_dict

class CkptExcludeKernel:
    """
    Removes the keys from the given model state_dict that match the key_pattern.

    Args:
        key_pattern: Patterns used to select the keys in the state_dict
            that are eligible for this kernel.
    """

    def __init__(self, key_pattern: List[str]):
        self.key_pattern = key_pattern

    def __call__(self, state_dict: Dict):
        """
        Args:
            state_dict: A dictionary representing the given checkpoint's state dict.
        """

        exclude_keys = _unix_pattern_to_parameter_names(
            self.key_pattern, state_dict.keys()
        )
        include_keys = set(state_dict.keys()) - exclude_keys

        new_state_dict = {}
        for key in include_keys:
            new_state_dict[key] = state_dict[key]

        return new_state_dict

class CkptPrependKernel:
    """
    Prepends the given pattern to all the keys in the checkpoint state dict after
    selecting them with key_pattern.

    For instance, if prepend_pattern  = "some_prepend." and
    key_pattern = ["model.head"], this kernel would prepend "some_prepend." to
    "model.key", thus renaming the key "model.head" to "some_prepend.model.head".

    Args:
        prepend_pattern: The pattern to prepend the keys in the state_dict with.
        key_pattern: Patterns used to select the keys in the state_dict
            that are eligible for this kernel.
    """

    def __init__(self, prepend_pattern: str, key_pattern: List[str]):
        self.prepend_pattern = prepend_pattern
        self.key_pattern = key_pattern

    def __call__(self, state_dict: Dict):
        """
        Args:
            state_dict: A dictionary representing the given checkpoint's state dict.
        """

        all_keys = set(state_dict.keys())

        include_keys = set(state_dict.keys())
        if self.key_pattern is not None:
            include_keys = _unix_pattern_to_parameter_names(
                self.key_pattern, state_dict.keys()
            )

        excluded_keys = all_keys - include_keys

        # Add excluded keys from re-mapping
        new_state_dict = {}
        for k in excluded_keys:
            new_state_dict[k] = state_dict[k]

        # Add keys from remapping
        for key in include_keys:
            new_state_dict[self.prepend_pattern + key] = state_dict[key]

        return new_state_dict

class CkptRenameWithCopyKernel:
    """
    Renames and also optionally creates copyies of the key-value pairs in the checkpoint
    state dict. Before doing so, selects the keys to which to apply this kernel by
    using key_pattern.

    For instance, if source_pattern  = "model.head" and
    target_patterns = ["model.head_1", "model.head_2"], this kernel would
    rename the key "model.head" to "model.head_1" and will also create a copy of the
    "model.head" and assign it a new name "model.head_2".

    Args:
        source_pattern: The pattern that needs to be renamed in the current
            checkpoint state_dict.
        target_patterns: A list of patterns to which the source_pattern is to be
            renamed to it. If the list has more than one element, it creates multiple
            copies of the source_pattern value and assigns then the names given in
            target_pattern.
        key_pattern: Patterns used to select the keys in the state_dict
            that are eligible for this kernel.
    """

    def __init__(
        self,
        source_pattern: str,
        target_patterns: List[str],
        key_pattern: Optional[List[str]] = None,
    ):
        self.source_pattern = source_pattern
        self.target_patterns = target_patterns
        self.key_pattern = key_pattern

    def __call__(self, state_dict: Dict):
        """
        Args:
            state_dict: A dictionary representing the given checkpoint's state dict.
        """

        # Replaces only first occurences

        all_keys = set(state_dict.keys())

        include_keys = set(state_dict.keys())
        if self.key_pattern is not None:
            include_keys = _unix_pattern_to_parameter_names(
                self.key_pattern, state_dict.keys()
            )

        excluded_keys = all_keys - include_keys

        # Add excluded keys from re-mapping
        new_state_dict = {}
        for k in excluded_keys:
            new_state_dict[k] = state_dict[k]

        # Add keys from remapping
        for key in include_keys:
            if self.source_pattern in key:
                for target_pattern in self.target_patterns:
                    new_key = key.replace(self.source_pattern, target_pattern, 1)
                    new_state_dict[new_key] = state_dict[key]
            else:
                new_state_dict[key] = state_dict[key]

        return new_state_dict


Prompt:
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import fnmatch
import logging
from typing import Any, Callable, Dict, List, Optional, Set, Type, Union

import hydra
import torch
import torch.nn as nn
from iopath.common.file_io import g_pathmgr
from omegaconf import OmegaConf

from .model_wrappers import MIMOHeadWrapper

def _unix_pattern_to_parameter_names(
    constraints: List[str], all_parameter_names: Set[str]
) -> Union[None, Set[str]]:

    parameter_names = []
    for param_name in constraints:
        matching_parameters = set(fnmatch.filter(all_parameter_names, param_name))
        assert (
            len(matching_parameters) > 0
        ), f"param_names {param_name} don't match any param in the given names."
        parameter_names.append(matching_parameters)
    return set.union(*parameter_names)

class CkptIncludeKernel:
    """
    Includes only the keys from the given model state_dict that match the key_pattern.
    Rest of the keys are removed from the given state_dict.

    Args:
        key_pattern: Patterns used to select the keys in the state_dict
            that are eligible for this kernel.
    """

    def __init__(self, key_pattern: List[str]):
        self.key_pattern = key_pattern

    def __call__(self, state_dict: Dict):
        """
        Args:
            state_dict: A dictionary representing the given checkpoint's state dict.
        """

        include_keys = _unix_pattern_to_parameter_names(
            self.key_pattern, state_dict.keys()
        )

        new_state_dict = {}
        for key in include_keys:
            new_state_dict[key] = state_dict[key]

        return new_state_dict

class CkptExcludeKernel:
    """
    Removes the keys from the given model state_dict that match the key_pattern.

    Args:
        key_pattern: Patterns used to select the keys in the state_dict
            that are eligible for this kernel.
    """

    def __init__(self, key_pattern: List[str]):
        self.key_pattern = key_pattern

    def __call__(self, state_dict: Dict):
        """
        Args:
            state_dict: A dictionary representing the given checkpoint's state dict.
        """

        exclude_keys = _unix_pattern_to_parameter_names(
            self.key_pattern, state_dict.keys()
        )
        include_keys = set(state_dict.keys()) - exclude_keys

        new_state_dict = {}
        for key in include_keys:
            new_state_dict[key] = state_dict[key]

        return new_state_dict

class CkptPrependKernel:
    """
    Prepends the given pattern to all the keys in the checkpoint state dict after
    selecting them with key_pattern.

    For instance, if prepend_pattern  = "some_prepend." and
    key_pattern = ["model.head"], this kernel would prepend "some_prepend." to
    "model.key", thus renaming the key "model.head" to "some_prepend.model.head".

    Args:
        prepend_pattern: The pattern to prepend the keys in the state_dict with.
        key_pattern: Patterns used to select the keys in the state_dict
            that are eligible for this kernel.
    """

    def __init__(self, prepend_pattern: str, key_pattern: List[str]):
        self.prepend_pattern = prepend_pattern
        self.key_pattern = key_pattern

    def __call__(self, state_dict: Dict):
        """
        Args:
            state_dict: A dictionary representing the given checkpoint's state dict.
        """

        all_keys = set(state_dict.keys())

        include_keys = set(state_dict.keys())
        if self.key_pattern is not None:
            include_keys = _unix_pattern_to_parameter_names(
                self.key_pattern, state_dict.keys()
            )

        excluded_keys = all_keys - include_keys

        # Add excluded keys from re-mapping
        new_state_dict = {}
        for k in excluded_keys:
            new_state_dict[k] = state_dict[k]

        # Add keys from remapping
        for key in include_keys:
            new_state_dict[self.prepend_pattern + key] = state_dict[key]

        return new_state_dict

class CkptRenameWithCopyKernel:
    """
    Renames and also optionally creates copyies of the key-value pairs in the checkpoint
    state dict. Before doing so, selects the keys to which to apply this kernel by
    using key_pattern.

    For instance, if source_pattern  = "model.head" and
    target_patterns = ["model.head_1", "model.head_2"], this kernel would
    rename the key "model.head" to "model.head_1" and will also create a copy of the
    "model.head" and assign it a new name "model.head_2".

    Args:
        source_pattern: The pattern that needs to be renamed in the current
            checkpoint state_dict.
        target_patterns: A list of patterns to which the source_pattern is to be
            renamed to it. If the list has more than one element, it creates multiple
            copies of the source_pattern value and assigns then the names given in
            target_pattern.
        key_pattern: Patterns used to select the keys in the state_dict
            that are eligible for this kernel.
    """

    def __init__(
        self,
        source_pattern: str,
        target_patterns: List[str],
        key_pattern: Optional[List[str]] = None,
    ):
        self.source_pattern = source_pattern
        self.target_patterns = target_patterns
        self.key_pattern = key_pattern

    def __call__(self, state_dict: Dict):
        """
        Args:
            state_dict: A dictionary representing the given checkpoint's state dict.
        """

        # Replaces only first occurences
        all_keys = set(state_dict.keys())

        include_keys = set(state_dict.keys())
        if self.key_pattern is not None:
            include_keys = _unix_pattern_to_parameter_names(
                self.key_pattern, state_dict.keys()
            )

        excluded_keys = all_keys - include_keys

        # Add excluded keys from re-mapping
        new_state_dict = {}
        for k in excluded_keys:
            new_state_dict[k] = state_dict[k]

        # Add keys from remapping
        for key in include_keys:
            if self.source_pattern in key:
                for target_pattern in self.target_patterns:
                    new_key = key.replace(self.source_pattern, target_pattern, 1)
                    new_state_dict[new_key] = state_dict[key]
            else:
                new_state_dict[key] = state_dict[key]

        return new_state_dict

def load_checkpoint(
    path_list: List[str],
    pick_recursive_keys: Optional[List[str]] = None,
    map_location: str = "cpu",
) -> Any:
    """
    Loads a checkpoint from the specified path.

    Args:
        path_list: A list of paths which contain the checkpoint. Each element
            is tried (in order) until a file that exists is found. That file is then
            used to read the checkpoint.
        pick_recursive_keys: Picks sub dicts from the loaded checkpoint if not None.
            For pick_recursive_keys = ["a", "b"], will return checkpoint_dict["a"]["b"]
        map_location (str): a function, torch.device, string or a dict specifying how to
            remap storage locations

    Returns: Model with the matchin pre-trained weights loaded.
    """
    path_exists = False
    for path in path_list:
        if g_pathmgr.exists(path):
            path_exists = True
            break

    if not path_exists:
        raise ValueError(f"No path exists in {path_list}")

    with g_pathmgr.open(path, "rb") as f:
        checkpoint = torch.load(f, map_location=map_location)

    logging.info(f"Loaded checkpoint from {path}")
    if pick_recursive_keys is not None:
        for key in pick_recursive_keys:
            checkpoint = checkpoint[key]
    return checkpoint

def load_checkpoint_and_apply_kernels(
    checkpoint_path: str,
    checkpoint_kernels: List[Callable] = None,
    ckpt_state_dict_key: str = "state_dict",
    map_location: str = None,
) -> nn.Module:
    """
    Performs checkpoint loading with a variety of pre-processing kernel applied in
    sequence.

    Args:
        checkpoint_path (str): Path to the checkpoint.
        checkpoint_kernels List(Callable): A list of checkpoint processing kernels
            to apply in the specified order. Supported kernels include `CkptIncludeKernel`,
            `CkptExcludeKernel`, etc. These kernels are applied in the
            given order.
        ckpt_state_dict_key (str): Key containing the model state dict.
        map_location (str): a function, torch.device, string or a dict specifying how to
            remap storage locations

    Returns: Model with the matchin pre-trained weights loaded.
    """

    assert g_pathmgr.exists(checkpoint_path), "Checkpoint '{}' not found".format(
        checkpoint_path
    )

    # Load the checkpoint on CPU to avoid GPU mem spike.
    with g_pathmgr.open(checkpoint_path, "rb") as f:
        checkpoint = torch.load(f, map_location=map_location)

    pre_train_dict = (
        checkpoint[ckpt_state_dict_key] if ckpt_state_dict_key else checkpoint
    )

    logging.info(
        "Loaded Checkpoint State Dict pre-kernel application: %s"
        % str(", ".join(list(pre_train_dict.keys())))
    )
    # Apply kernels
    if checkpoint_kernels is not None:
        for f in checkpoint_kernels:
            pre_train_dict = f(state_dict=pre_train_dict)

    logging.info(
        "Loaded Checkpoint State Dict Post-kernel application %s"
        % str(", ".join(list(pre_train_dict.keys())))
    )

    return pre_train_dict


Prompt:
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import fnmatch
import logging
from typing import Any, Callable, Dict, List, Optional, Set, Type, Union

import hydra
import torch
import torch.nn as nn
from iopath.common.file_io import g_pathmgr
from omegaconf import OmegaConf

from .model_wrappers import MIMOHeadWrapper

def _unix_pattern_to_parameter_names(
    constraints: List[str], all_parameter_names: Set[str]
) -> Union[None, Set[str]]:

    parameter_names = []
    for param_name in constraints:
        matching_parameters = set(fnmatch.filter(all_parameter_names, param_name))
        assert (
            len(matching_parameters) > 0
        ), f"param_names {param_name} don't match any param in the given names."
        parameter_names.append(matching_parameters)
    return set.union(*parameter_names)

class CkptIncludeKernel:
    """
    Includes only the keys from the given model state_dict that match the key_pattern.
    Rest of the keys are removed from the given state_dict.

    Args:
        key_pattern: Patterns used to select the keys in the state_dict
            that are eligible for this kernel.
    """

    def __init__(self, key_pattern: List[str]):
        self.key_pattern = key_pattern

    def __call__(self, state_dict: Dict):
        """
        Args:
            state_dict: A dictionary representing the given checkpoint's state dict.
        """

        include_keys = _unix_pattern_to_parameter_names(
            self.key_pattern, state_dict.keys()
        )

        new_state_dict = {}
        for key in include_keys:
            new_state_dict[key] = state_dict[key]

        return new_state_dict

class CkptExcludeKernel:
    """
    Removes the keys from the given model state_dict that match the key_pattern.

    Args:
        key_pattern: Patterns used to select the keys in the state_dict
            that are eligible for this kernel.
    """

    def __init__(self, key_pattern: List[str]):
        self.key_pattern = key_pattern

    def __call__(self, state_dict: Dict):
        """
        Args:
            state_dict: A dictionary representing the given checkpoint's state dict.
        """

        exclude_keys = _unix_pattern_to_parameter_names(
            self.key_pattern, state_dict.keys()
        )
        include_keys = set(state_dict.keys()) - exclude_keys

        new_state_dict = {}
        for key in include_keys:
            new_state_dict[key] = state_dict[key]

        return new_state_dict

class CkptPrependKernel:
    """
    Prepends the given pattern to all the keys in the checkpoint state dict after
    selecting them with key_pattern.

    For instance, if prepend_pattern  = "some_prepend." and
    key_pattern = ["model.head"], this kernel would prepend "some_prepend." to
    "model.key", thus renaming the key "model.head" to "some_prepend.model.head".

    Args:
        prepend_pattern: The pattern to prepend the keys in the state_dict with.
        key_pattern: Patterns used to select the keys in the state_dict
            that are eligible for this kernel.
    """

    def __init__(self, prepend_pattern: str, key_pattern: List[str]):
        self.prepend_pattern = prepend_pattern
        self.key_pattern = key_pattern

    def __call__(self, state_dict: Dict):
        """
        Args:
            state_dict: A dictionary representing the given checkpoint's state dict.
        """

        all_keys = set(state_dict.keys())

        include_keys = set(state_dict.keys())
        if self.key_pattern is not None:
            include_keys = _unix_pattern_to_parameter_names(
                self.key_pattern, state_dict.keys()
            )

        excluded_keys = all_keys - include_keys

        # Add excluded keys from re-mapping
        new_state_dict = {}
        for k in excluded_keys:
            new_state_dict[k] = state_dict[k]

        # Add keys from remapping
        for key in include_keys:
            new_state_dict[self.prepend_pattern + key] = state_dict[key]

        return new_state_dict

class CkptRenameWithCopyKernel:
    """
    Renames and also optionally creates copyies of the key-value pairs in the checkpoint
    state dict. Before doing so, selects the keys to which to apply this kernel by
    using key_pattern.

    For instance, if source_pattern  = "model.head" and
    target_patterns = ["model.head_1", "model.head_2"], this kernel would
    rename the key "model.head" to "model.head_1" and will also create a copy of the
    "model.head" and assign it a new name "model.head_2".

    Args:
        source_pattern: The pattern that needs to be renamed in the current
            checkpoint state_dict.
        target_patterns: A list of patterns to which the source_pattern is to be
            renamed to it. If the list has more than one element, it creates multiple
            copies of the source_pattern value and assigns then the names given in
            target_pattern.
        key_pattern: Patterns used to select the keys in the state_dict
            that are eligible for this kernel.
    """

    def __init__(
        self,
        source_pattern: str,
        target_patterns: List[str],
        key_pattern: Optional[List[str]] = None,
    ):
        self.source_pattern = source_pattern
        self.target_patterns = target_patterns
        self.key_pattern = key_pattern

    def __call__(self, state_dict: Dict):
        """
        Args:
            state_dict: A dictionary representing the given checkpoint's state dict.
        """

        # Replaces only first occurences
        all_keys = set(state_dict.keys())

        include_keys = set(state_dict.keys())
        if self.key_pattern is not None:
            include_keys = _unix_pattern_to_parameter_names(
                self.key_pattern, state_dict.keys()
            )

        excluded_keys = all_keys - include_keys

        # Add excluded keys from re-mapping
        new_state_dict = {}
        for k in excluded_keys:
            new_state_dict[k] = state_dict[k]

        # Add keys from remapping
        for key in include_keys:
            if self.source_pattern in key:
                for target_pattern in self.target_patterns:
                    new_key = key.replace(self.source_pattern, target_pattern, 1)
                    new_state_dict[new_key] = state_dict[key]
            else:
                new_state_dict[key] = state_dict[key]

        return new_state_dict

def load_checkpoint(
    path_list: List[str],
    pick_recursive_keys: Optional[List[str]] = None,
    map_location: str = "cpu",
) -> Any:
    """
    Loads a checkpoint from the specified path.

    Args:
        path_list: A list of paths which contain the checkpoint. Each element
            is tried (in order) until a file that exists is found. That file is then
            used to read the checkpoint.
        pick_recursive_keys: Picks sub dicts from the loaded checkpoint if not None.
            For pick_recursive_keys = ["a", "b"], will return checkpoint_dict["a"]["b"]
        map_location (str): a function, torch.device, string or a dict specifying how to
            remap storage locations

    Returns: Model with the matchin pre-trained weights loaded.
    """
    path_exists = False
    for path in path_list:
        if g_pathmgr.exists(path):
            path_exists = True
            break

    if not path_exists:
        raise ValueError(f"No path exists in {path_list}")

    with g_pathmgr.open(path, "rb") as f:
        checkpoint = torch.load(f, map_location=map_location)

    logging.info(f"Loaded checkpoint from {path}")
    if pick_recursive_keys is not None:
        for key in pick_recursive_keys:
            checkpoint = checkpoint[key]
    return checkpoint

def load_checkpoint_and_apply_kernels(
    checkpoint_path: str,
    checkpoint_kernels: List[Callable] = None,
    ckpt_state_dict_key: str = "state_dict",
    map_location: str = None,
) -> nn.Module:
    """
    Performs checkpoint loading with a variety of pre-processing kernel applied in
    sequence.

    Args:
        checkpoint_path (str): Path to the checkpoint.
        checkpoint_kernels List(Callable): A list of checkpoint processing kernels
            to apply in the specified order. Supported kernels include `CkptIncludeKernel`,
            `CkptExcludeKernel`, etc. These kernels are applied in the
            given order.
        ckpt_state_dict_key (str): Key containing the model state dict.
        map_location (str): a function, torch.device, string or a dict specifying how to
            remap storage locations

    Returns: Model with the matchin pre-trained weights loaded.
    """
    assert g_pathmgr.exists(checkpoint_path), "Checkpoint '{}' not found".format(
        checkpoint_path
    )

    # Load the checkpoint on CPU to avoid GPU mem spike.
    with g_pathmgr.open(checkpoint_path, "rb") as f:
        checkpoint = torch.load(f, map_location=map_location)

    pre_train_dict = (
        checkpoint[ckpt_state_dict_key] if ckpt_state_dict_key else checkpoint
    )

    logging.info(
        "Loaded Checkpoint State Dict pre-kernel application: %s"
        % str(", ".join(list(pre_train_dict.keys())))
    )
    # Apply kernels
    if checkpoint_kernels is not None:
        for f in checkpoint_kernels:
            pre_train_dict = f(state_dict=pre_train_dict)

    logging.info(
        "Loaded Checkpoint State Dict Post-kernel application %s"
        % str(", ".join(list(pre_train_dict.keys())))
    )

    return pre_train_dict

def load_state_dict_into_model(state_dict: Dict, model: nn.Module, strict: bool = True):
    """
    Loads a state dict into the given model.

    Args:
        state_dict: A dictionary containing the model's
            state dict, or a subset if strict is False
        model: Model to load the checkpoint weights into
        strict: raise if the state_dict has missing state keys
    """

    missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False)
    err = "State key mismatch."
    if unexpected_keys:
        err += f" Unexpected keys: {unexpected_keys}."
    if missing_keys:
        err += f" Missing keys: {missing_keys}."
    if unexpected_keys or missing_keys:
        if not unexpected_keys and not strict:
            logging.warning(err)
        else:
            raise KeyError(err)
    return model


Prompt:
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import copy
from dataclasses import dataclass, field
from typing import Dict, List, Mapping, Optional, Sequence

import numpy as np

import torch
import torch.nn as nn
from omnivision.data.api import VisionSample

class MIMOHeadWrapper(nn.Module):
    """Attaches multiple input multiple output heads to the trunk using forward hooks.

    Args:
        trunk: Any model to which you want to attach the heads to.
        heads: A list of dicts with the following keys:
            fork_module: The module which the head will be applied to. It can be an
                empty string, in which case the head is attached to the trunk's output.
            head: The head which is to be attached.
            input_key: The head will only run on inputs with this key. If set to
                `None` the head will be applied to all inputs.
            output_key: The head will produce this output key. If set to `None`, the
                output key will be the same as the input key.

            An example heads value can look like -
            ```
            [
                {
                    "fork_module": "layer_1.layer_a.layer_alpha",
                    "head": nn.Linear(in_feat, out_feat),
                    "input_key": "dataset_1",
                    "output_key": "out_1",
                },
                {
                    "fork_module": "",
                    "head": nn.Linear(in_feat, out_feat),
                    "input_key": "dataset_1",
                    "output_key": "out_2",
                },
                {
                    "fork_module": "",
                    "head": nn.Linear(in_feat, out_feat),
                    "input_key": "dataset_2",
                    "output_key": "out_3",
                },
                {
                    "fork_module": "",
                    "head": nn.Conv2d(in_feat, out_feat),
                    "input_key": None,
                    "output_key": None,
                },
            ]
            ```
        trunk_fields: A list of dicts with the following keys:
            input_key: The input key this rule applies to. If `None`, applies to all
                inputs.
            args: These specific keys will be fetched from the sample and passed as
                *args to the trunk for the specified `input_key`.
            kwargs: These specific keys will be fetched from the sample and passed as
                **kwargs to the trunk for the specified `input_key`.

            Example -
            ```
            [
                {
                    "input_key": "dataset_1",
                    "args": ["vision"]
                },
                {
                    "input_key": "dataset_2",
                    "args": ["vision"],
                    "kwargs": {"mask": "mask"}
                },
            ]
            ```

        Note that two heads cannot produce the same output key in the same forward pass.

    Returns:
        A dict with keys corresponding to the output keys which match with the input key.
    """

    @dataclass
    class HeadArgs:
        fork_module: str
        head: nn.Module
        input_key: Optional[str]
        output_key: Optional[str]

    @dataclass
    class TrunkFieldArgs:
        input_key: Optional[str]
        args: List[str] = field(default_factory=list)
        kwargs: Dict[str, str] = field(default_factory=dict)

    def __init__(
        self,
        trunk: nn.Module,
        heads: List[Dict],
        trunk_fields: List[Dict],
        handle_list_inputs=False,
    ) -> None:
        """WARNING: handle_list_inputs is a hack which needs to be refactored away."""
        super().__init__()

        self.trunk = trunk
        self.handle_list_inputs = handle_list_inputs

        # cast to HeadArgs for input validation
        heads = [self.HeadArgs(**head_dict) for head_dict in heads]
        # cast to TrunkFieldArgs for input validation
        trunk_fields = [
            self.TrunkFieldArgs(**trunk_fields_dict)
            for trunk_fields_dict in trunk_fields
        ]

        self.head_name_to_fork_module = {}
        self.heads = nn.ModuleList()
        self.head_input_keys = []
        self.head_output_keys = []
        self.head_fork_modules = []

        for head_args in heads:
            self.heads.append(head_args.head)
            self.head_input_keys.append(head_args.input_key)
            self.head_output_keys.append(head_args.output_key)
            self.head_fork_modules.append(head_args.fork_module)

        self.trunk_field_args = {}
        self.trunk_field_kwargs = {}
        for trunk_fields_elem in trunk_fields:
            input_key = trunk_fields_elem.input_key
            if input_key in self.trunk_field_args:
                raise KeyError(
                    f"Multiple trunk_fields specified for the same input_key: {input_key}"
                )
            self.trunk_field_args[input_key] = trunk_fields_elem.args
            self.trunk_field_kwargs[input_key] = trunk_fields_elem.kwargs

        # outputs is used as a temporary storage of the head outputs
        self.outputs = {}

        # input_key is used to specify which key is currently being processed
        self.input_key = None

        # handles to the hooks which can be used for removing the hooks if needed
        self.hook_handles = []
        self._register_hooks()

    def _register_hooks(self):
        for i, head in enumerate(self.heads):
            fork_module_name = self.head_fork_modules[i]

            def hook_fn(
                module,
                module_in,
                module_out,
                # the following variables are passed as kwargs in the closure to avoid
                # late binding in python
                head_method=head,
                in_key=self.head_input_keys[i],
                out_key=self.head_output_keys[i],
            ):

                if in_key is not None and self.input_key != in_key:
                    return
                if out_key is None:
                    out_key = self.input_key
                if out_key in self.outputs:
                    # reset state before raising
                    self.outputs = {}
                    self.input_key = None
                    raise ValueError(
                        f"Two heads produced the same output key `{out_key}` during forward"
                    )
                self.outputs[out_key] = head_method(module_out)


Prompt:
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import copy
from dataclasses import dataclass, field
from typing import Dict, List, Mapping, Optional, Sequence

import numpy as np

import torch
import torch.nn as nn
from omnivision.data.api import VisionSample

class MIMOHeadWrapper(nn.Module):
    """Attaches multiple input multiple output heads to the trunk using forward hooks.

    Args:
        trunk: Any model to which you want to attach the heads to.
        heads: A list of dicts with the following keys:
            fork_module: The module which the head will be applied to. It can be an
                empty string, in which case the head is attached to the trunk's output.
            head: The head which is to be attached.
            input_key: The head will only run on inputs with this key. If set to
                `None` the head will be applied to all inputs.
            output_key: The head will produce this output key. If set to `None`, the
                output key will be the same as the input key.

            An example heads value can look like -
            ```
            [
                {
                    "fork_module": "layer_1.layer_a.layer_alpha",
                    "head": nn.Linear(in_feat, out_feat),
                    "input_key": "dataset_1",
                    "output_key": "out_1",
                },
                {
                    "fork_module": "",
                    "head": nn.Linear(in_feat, out_feat),
                    "input_key": "dataset_1",
                    "output_key": "out_2",
                },
                {
                    "fork_module": "",
                    "head": nn.Linear(in_feat, out_feat),
                    "input_key": "dataset_2",
                    "output_key": "out_3",
                },
                {
                    "fork_module": "",
                    "head": nn.Conv2d(in_feat, out_feat),
                    "input_key": None,
                    "output_key": None,
                },
            ]
            ```
        trunk_fields: A list of dicts with the following keys:
            input_key: The input key this rule applies to. If `None`, applies to all
                inputs.
            args: These specific keys will be fetched from the sample and passed as
                *args to the trunk for the specified `input_key`.
            kwargs: These specific keys will be fetched from the sample and passed as
                **kwargs to the trunk for the specified `input_key`.

            Example -
            ```
            [
                {
                    "input_key": "dataset_1",
                    "args": ["vision"]
                },
                {
                    "input_key": "dataset_2",
                    "args": ["vision"],
                    "kwargs": {"mask": "mask"}
                },
            ]
            ```

        Note that two heads cannot produce the same output key in the same forward pass.

    Returns:
        A dict with keys corresponding to the output keys which match with the input key.
    """

    @dataclass
    class HeadArgs:
        fork_module: str
        head: nn.Module
        input_key: Optional[str]
        output_key: Optional[str]

    @dataclass
    class TrunkFieldArgs:
        input_key: Optional[str]
        args: List[str] = field(default_factory=list)
        kwargs: Dict[str, str] = field(default_factory=dict)

    def __init__(
        self,
        trunk: nn.Module,
        heads: List[Dict],
        trunk_fields: List[Dict],
        handle_list_inputs=False,
    ) -> None:
        """WARNING: handle_list_inputs is a hack which needs to be refactored away."""
        super().__init__()

        self.trunk = trunk
        self.handle_list_inputs = handle_list_inputs

        # cast to HeadArgs for input validation
        heads = [self.HeadArgs(**head_dict) for head_dict in heads]
        # cast to TrunkFieldArgs for input validation
        trunk_fields = [
            self.TrunkFieldArgs(**trunk_fields_dict)
            for trunk_fields_dict in trunk_fields
        ]

        self.head_name_to_fork_module = {}
        self.heads = nn.ModuleList()
        self.head_input_keys = []
        self.head_output_keys = []
        self.head_fork_modules = []

        for head_args in heads:
            self.heads.append(head_args.head)
            self.head_input_keys.append(head_args.input_key)
            self.head_output_keys.append(head_args.output_key)
            self.head_fork_modules.append(head_args.fork_module)

        self.trunk_field_args = {}
        self.trunk_field_kwargs = {}
        for trunk_fields_elem in trunk_fields:
            input_key = trunk_fields_elem.input_key
            if input_key in self.trunk_field_args:
                raise KeyError(
                    f"Multiple trunk_fields specified for the same input_key: {input_key}"
                )
            self.trunk_field_args[input_key] = trunk_fields_elem.args
            self.trunk_field_kwargs[input_key] = trunk_fields_elem.kwargs

        # outputs is used as a temporary storage of the head outputs
        self.outputs = {}

        # input_key is used to specify which key is currently being processed
        self.input_key = None

        # handles to the hooks which can be used for removing the hooks if needed
        self.hook_handles = []
        self._register_hooks()

    def _register_hooks(self):
        for i, head in enumerate(self.heads):
            fork_module_name = self.head_fork_modules[i]

            def hook_fn(
                module,
                module_in,
                module_out,
                # the following variables are passed as kwargs in the closure to avoid
                # late binding in python
                head_method=head,
                in_key=self.head_input_keys[i],
                out_key=self.head_output_keys[i],
            ):
                if in_key is not None and self.input_key != in_key:
                    return
                if out_key is None:
                    out_key = self.input_key
                if out_key in self.outputs:
                    # reset state before raising
                    self.outputs = {}
                    self.input_key = None
                    raise ValueError(
                        f"Two heads produced the same output key `{out_key}` during forward"
                    )
                self.outputs[out_key] = head_method(module_out)

            fork_module = self.trunk.get_submodule(fork_module_name)
            self.hook_handles.append(fork_module.register_forward_hook(hook_fn))

    def _get_trunk_fields(self):

        fields_args = self.trunk_field_args.get(self.input_key)
        fields_kwargs = self.trunk_field_kwargs.get(self.input_key)
        if fields_args is None:
            assert fields_kwargs is None
            fields_args = self.trunk_field_args.get(None)
            fields_kwargs = self.trunk_field_kwargs.get(None)
            if fields_args is None:
                assert fields_kwargs is None
                raise ValueError(
                    f"No trunk fields specified for input key: {self.input_key}"
                )
        return fields_args, fields_kwargs


Prompt:
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import copy
from dataclasses import dataclass, field
from typing import Dict, List, Mapping, Optional, Sequence

import numpy as np

import torch
import torch.nn as nn
from omnivision.data.api import VisionSample

class MIMOHeadWrapper(nn.Module):
    """Attaches multiple input multiple output heads to the trunk using forward hooks.

    Args:
        trunk: Any model to which you want to attach the heads to.
        heads: A list of dicts with the following keys:
            fork_module: The module which the head will be applied to. It can be an
                empty string, in which case the head is attached to the trunk's output.
            head: The head which is to be attached.
            input_key: The head will only run on inputs with this key. If set to
                `None` the head will be applied to all inputs.
            output_key: The head will produce this output key. If set to `None`, the
                output key will be the same as the input key.

            An example heads value can look like -
            ```
            [
                {
                    "fork_module": "layer_1.layer_a.layer_alpha",
                    "head": nn.Linear(in_feat, out_feat),
                    "input_key": "dataset_1",
                    "output_key": "out_1",
                },
                {
                    "fork_module": "",
                    "head": nn.Linear(in_feat, out_feat),
                    "input_key": "dataset_1",
                    "output_key": "out_2",
                },
                {
                    "fork_module": "",
                    "head": nn.Linear(in_feat, out_feat),
                    "input_key": "dataset_2",
                    "output_key": "out_3",
                },
                {
                    "fork_module": "",
                    "head": nn.Conv2d(in_feat, out_feat),
                    "input_key": None,
                    "output_key": None,
                },
            ]
            ```
        trunk_fields: A list of dicts with the following keys:
            input_key: The input key this rule applies to. If `None`, applies to all
                inputs.
            args: These specific keys will be fetched from the sample and passed as
                *args to the trunk for the specified `input_key`.
            kwargs: These specific keys will be fetched from the sample and passed as
                **kwargs to the trunk for the specified `input_key`.

            Example -
            ```
            [
                {
                    "input_key": "dataset_1",
                    "args": ["vision"]
                },
                {
                    "input_key": "dataset_2",
                    "args": ["vision"],
                    "kwargs": {"mask": "mask"}
                },
            ]
            ```

        Note that two heads cannot produce the same output key in the same forward pass.

    Returns:
        A dict with keys corresponding to the output keys which match with the input key.
    """

    @dataclass
    class HeadArgs:
        fork_module: str
        head: nn.Module
        input_key: Optional[str]
        output_key: Optional[str]

    @dataclass
    class TrunkFieldArgs:
        input_key: Optional[str]
        args: List[str] = field(default_factory=list)
        kwargs: Dict[str, str] = field(default_factory=dict)

    def __init__(
        self,
        trunk: nn.Module,
        heads: List[Dict],
        trunk_fields: List[Dict],
        handle_list_inputs=False,
    ) -> None:
        """WARNING: handle_list_inputs is a hack which needs to be refactored away."""
        super().__init__()

        self.trunk = trunk
        self.handle_list_inputs = handle_list_inputs

        # cast to HeadArgs for input validation
        heads = [self.HeadArgs(**head_dict) for head_dict in heads]
        # cast to TrunkFieldArgs for input validation
        trunk_fields = [
            self.TrunkFieldArgs(**trunk_fields_dict)
            for trunk_fields_dict in trunk_fields
        ]

        self.head_name_to_fork_module = {}
        self.heads = nn.ModuleList()
        self.head_input_keys = []
        self.head_output_keys = []
        self.head_fork_modules = []

        for head_args in heads:
            self.heads.append(head_args.head)
            self.head_input_keys.append(head_args.input_key)
            self.head_output_keys.append(head_args.output_key)
            self.head_fork_modules.append(head_args.fork_module)

        self.trunk_field_args = {}
        self.trunk_field_kwargs = {}
        for trunk_fields_elem in trunk_fields:
            input_key = trunk_fields_elem.input_key
            if input_key in self.trunk_field_args:
                raise KeyError(
                    f"Multiple trunk_fields specified for the same input_key: {input_key}"
                )
            self.trunk_field_args[input_key] = trunk_fields_elem.args
            self.trunk_field_kwargs[input_key] = trunk_fields_elem.kwargs

        # outputs is used as a temporary storage of the head outputs
        self.outputs = {}

        # input_key is used to specify which key is currently being processed
        self.input_key = None

        # handles to the hooks which can be used for removing the hooks if needed
        self.hook_handles = []
        self._register_hooks()

    def _register_hooks(self):
        for i, head in enumerate(self.heads):
            fork_module_name = self.head_fork_modules[i]

            def hook_fn(
                module,
                module_in,
                module_out,
                # the following variables are passed as kwargs in the closure to avoid
                # late binding in python
                head_method=head,
                in_key=self.head_input_keys[i],
                out_key=self.head_output_keys[i],
            ):
                if in_key is not None and self.input_key != in_key:
                    return
                if out_key is None:
                    out_key = self.input_key
                if out_key in self.outputs:
                    # reset state before raising
                    self.outputs = {}
                    self.input_key = None
                    raise ValueError(
                        f"Two heads produced the same output key `{out_key}` during forward"
                    )
                self.outputs[out_key] = head_method(module_out)

            fork_module = self.trunk.get_submodule(fork_module_name)
            self.hook_handles.append(fork_module.register_forward_hook(hook_fn))

    def _get_trunk_fields(self):
        fields_args = self.trunk_field_args.get(self.input_key)
        fields_kwargs = self.trunk_field_kwargs.get(self.input_key)
        if fields_args is None:
            assert fields_kwargs is None
            fields_args = self.trunk_field_args.get(None)
            fields_kwargs = self.trunk_field_kwargs.get(None)
            if fields_args is None:
                assert fields_kwargs is None
                raise ValueError(
                    f"No trunk fields specified for input key: {self.input_key}"
                )
        return fields_args, fields_kwargs

    def forward_sub_batch(self, sub_batch, *args, **kwargs):

        assert isinstance(sub_batch, VisionSample), f"Received {type(sub_batch)}"
        fields_args, fields_kwargs = self._get_trunk_fields()
        sample_args = [getattr(sub_batch, arg) for arg in fields_args]
        sample_kwargs = {
            key: getattr(sub_batch, field) for key, field in fields_kwargs.items()
        }
        self.trunk(*sample_args, *args, **sample_kwargs, **kwargs)


Prompt:
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import copy
from dataclasses import dataclass, field
from typing import Dict, List, Mapping, Optional, Sequence

import numpy as np

import torch
import torch.nn as nn
from omnivision.data.api import VisionSample

class MIMOHeadWrapper(nn.Module):
    """Attaches multiple input multiple output heads to the trunk using forward hooks.

    Args:
        trunk: Any model to which you want to attach the heads to.
        heads: A list of dicts with the following keys:
            fork_module: The module which the head will be applied to. It can be an
                empty string, in which case the head is attached to the trunk's output.
            head: The head which is to be attached.
            input_key: The head will only run on inputs with this key. If set to
                `None` the head will be applied to all inputs.
            output_key: The head will produce this output key. If set to `None`, the
                output key will be the same as the input key.

            An example heads value can look like -
            ```
            [
                {
                    "fork_module": "layer_1.layer_a.layer_alpha",
                    "head": nn.Linear(in_feat, out_feat),
                    "input_key": "dataset_1",
                    "output_key": "out_1",
                },
                {
                    "fork_module": "",
                    "head": nn.Linear(in_feat, out_feat),
                    "input_key": "dataset_1",
                    "output_key": "out_2",
                },
                {
                    "fork_module": "",
                    "head": nn.Linear(in_feat, out_feat),
                    "input_key": "dataset_2",
                    "output_key": "out_3",
                },
                {
                    "fork_module": "",
                    "head": nn.Conv2d(in_feat, out_feat),
                    "input_key": None,
                    "output_key": None,
                },
            ]
            ```
        trunk_fields: A list of dicts with the following keys:
            input_key: The input key this rule applies to. If `None`, applies to all
                inputs.
            args: These specific keys will be fetched from the sample and passed as
                *args to the trunk for the specified `input_key`.
            kwargs: These specific keys will be fetched from the sample and passed as
                **kwargs to the trunk for the specified `input_key`.

            Example -
            ```
            [
                {
                    "input_key": "dataset_1",
                    "args": ["vision"]
                },
                {
                    "input_key": "dataset_2",
                    "args": ["vision"],
                    "kwargs": {"mask": "mask"}
                },
            ]
            ```

        Note that two heads cannot produce the same output key in the same forward pass.

    Returns:
        A dict with keys corresponding to the output keys which match with the input key.
    """

    @dataclass
    class HeadArgs:
        fork_module: str
        head: nn.Module
        input_key: Optional[str]
        output_key: Optional[str]

    @dataclass
    class TrunkFieldArgs:
        input_key: Optional[str]
        args: List[str] = field(default_factory=list)
        kwargs: Dict[str, str] = field(default_factory=dict)

    def __init__(
        self,
        trunk: nn.Module,
        heads: List[Dict],
        trunk_fields: List[Dict],
        handle_list_inputs=False,
    ) -> None:
        """WARNING: handle_list_inputs is a hack which needs to be refactored away."""
        super().__init__()

        self.trunk = trunk
        self.handle_list_inputs = handle_list_inputs

        # cast to HeadArgs for input validation
        heads = [self.HeadArgs(**head_dict) for head_dict in heads]
        # cast to TrunkFieldArgs for input validation
        trunk_fields = [
            self.TrunkFieldArgs(**trunk_fields_dict)
            for trunk_fields_dict in trunk_fields
        ]

        self.head_name_to_fork_module = {}
        self.heads = nn.ModuleList()
        self.head_input_keys = []
        self.head_output_keys = []
        self.head_fork_modules = []

        for head_args in heads:
            self.heads.append(head_args.head)
            self.head_input_keys.append(head_args.input_key)
            self.head_output_keys.append(head_args.output_key)
            self.head_fork_modules.append(head_args.fork_module)

        self.trunk_field_args = {}
        self.trunk_field_kwargs = {}
        for trunk_fields_elem in trunk_fields:
            input_key = trunk_fields_elem.input_key
            if input_key in self.trunk_field_args:
                raise KeyError(
                    f"Multiple trunk_fields specified for the same input_key: {input_key}"
                )
            self.trunk_field_args[input_key] = trunk_fields_elem.args
            self.trunk_field_kwargs[input_key] = trunk_fields_elem.kwargs

        # outputs is used as a temporary storage of the head outputs
        self.outputs = {}

        # input_key is used to specify which key is currently being processed
        self.input_key = None

        # handles to the hooks which can be used for removing the hooks if needed
        self.hook_handles = []
        self._register_hooks()

    def _register_hooks(self):
        for i, head in enumerate(self.heads):
            fork_module_name = self.head_fork_modules[i]

            def hook_fn(
                module,
                module_in,
                module_out,
                # the following variables are passed as kwargs in the closure to avoid
                # late binding in python
                head_method=head,
                in_key=self.head_input_keys[i],
                out_key=self.head_output_keys[i],
            ):
                if in_key is not None and self.input_key != in_key:
                    return
                if out_key is None:
                    out_key = self.input_key
                if out_key in self.outputs:
                    # reset state before raising
                    self.outputs = {}
                    self.input_key = None
                    raise ValueError(
                        f"Two heads produced the same output key `{out_key}` during forward"
                    )
                self.outputs[out_key] = head_method(module_out)

            fork_module = self.trunk.get_submodule(fork_module_name)
            self.hook_handles.append(fork_module.register_forward_hook(hook_fn))

    def _get_trunk_fields(self):
        fields_args = self.trunk_field_args.get(self.input_key)
        fields_kwargs = self.trunk_field_kwargs.get(self.input_key)
        if fields_args is None:
            assert fields_kwargs is None
            fields_args = self.trunk_field_args.get(None)
            fields_kwargs = self.trunk_field_kwargs.get(None)
            if fields_args is None:
                assert fields_kwargs is None
                raise ValueError(
                    f"No trunk fields specified for input key: {self.input_key}"
                )
        return fields_args, fields_kwargs

    def forward_sub_batch(self, sub_batch, *args, **kwargs):
        assert isinstance(sub_batch, VisionSample), f"Received {type(sub_batch)}"
        fields_args, fields_kwargs = self._get_trunk_fields()
        sample_args = [getattr(sub_batch, arg) for arg in fields_args]
        sample_kwargs = {
            key: getattr(sub_batch, field) for key, field in fields_kwargs.items()
        }
        self.trunk(*sample_args, *args, **sample_kwargs, **kwargs)

    def forward(self, batch, *args, **kwargs) -> Dict:

        assert isinstance(batch, Mapping)
        assert len(self.outputs) == 0
        for key, sub_batch in batch.items():
            self.input_key = key
            if self.handle_list_inputs and isinstance(sub_batch.vision, Sequence):
                # FIXME: this only handles list inputs for the field "vision"
                assert len(batch) == 1
                out_vals = []
                for e in sub_batch.vision:
                    e_batch = copy.copy(sub_batch)
                    e_batch.vision = e
                    self.forward_sub_batch(e_batch, *args, **kwargs)
                    assert len(self.outputs) == 1
                    out_key, out_val = self.outputs.popitem()
                    out_vals.append(out_val)
                return {out_key: torch.cat(out_vals)}
            else:
                self.forward_sub_batch(sub_batch, *args, **kwargs)
        outputs = self.outputs
        self.input_key = None
        self.outputs = {}
        return outputs


Prompt:
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

from dataclasses import dataclass, field, fields, is_dataclass, make_dataclass
from typing import Any, Callable, Dict

from torch.utils.data.dataloader import default_collate

@dataclass
class Batch:
    # the following are per batch args which are passed to the trainer
    # and are set to reasonable defaults
    model_fwd_kwargs: Dict = field(default_factory=dict)
    accum_steps: int = 1

def create_batch_sample_cls(cls):
    """Dynamically creates a dataclass which is a `Batch` and a `Sample`.

    This function also registers the class in globals() to make the class picklable.
    """

    cls_name = f"{Batch.__name__}{cls.__name__}"
    batch_sample_cls = make_dataclass(cls_name, fields=(), bases=(cls, Batch))
    batch_sample_cls.__module__ = __name__
    globals()[cls_name] = batch_sample_cls
    return cls


Prompt:
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import contextlib
import json
import logging
import math
import os
import sys
import time
from collections import OrderedDict
from dataclasses import dataclass
from typing import Any, Dict, List, Mapping, Optional, Sequence

import torch
import torch.distributed as dist
import torch.nn as nn
from hydra.utils import instantiate
from iopath.common.file_io import g_pathmgr
from omnivision.data.api import Sample
from omnivision.data.concat_dataset import ConcatDataset
from omnivision.data.torch_dataset import TorchDataset
from omnivision.losses import wrap_base_loss
from omnivision.optim import construct_optimizer
from omnivision.utils.train import (
    AverageMeter,
    copy_data_to_device,
    get_amp_type,
    get_machine_local_and_dist_rank,
    get_resume_checkpoint,
    is_dist_avail_and_initialized,
    makedir,
    ProgressMeter,
    set_seeds,
    setup_distributed_backend,
    setup_logging,
)

def chunk_batch_for_accum_steps(batch, accum_steps):
    return [get_chunk_from_data(batch, i, accum_steps) for i in range(accum_steps)]

def get_chunk_from_data(data, chunk_id, num_chunks):
    """
    Recursively splits all the tensors inside the passed data object into num_chunks.
    """
    if isinstance(data, torch.Tensor):
        assert len(data) % num_chunks == 0
        start = (len(data) // num_chunks) * chunk_id
        end = (len(data) // num_chunks) * (chunk_id + 1)
        return data[start:end]
    elif isinstance(data, Mapping):
        return {
            key: get_chunk_from_data(value, chunk_id, num_chunks)
            for key, value in data.items()
        }
    elif isinstance(data, Sequence):
        return [get_chunk_from_data(value, chunk_id, num_chunks) for value in data]
    elif isinstance(data, Sample):
        data_cls = type(data)
        data = data.__dict__
        return data_cls(**get_chunk_from_data(data, chunk_id, num_chunks))
    else:
        return data

@dataclass
class OmnivisionOptimAMPConf:
    enabled: bool = False
    amp_dtype: str = "float16"

@dataclass
class OmnivisionOptimConf:
    optimizer: torch.optim.Optimizer = None
    options: Optional[Dict[str, Any]] = None
    param_group_modifiers: Optional[List] = None
    amp: Optional[Dict[str, Any]] = None
    gradient_clip: Any = None

    def __post_init__(self):
        # amp

        if not isinstance(self.amp, OmnivisionOptimAMPConf):
            if self.amp is None:
                self.amp = {}
            assert isinstance(self.amp, Mapping)
            self.amp = OmnivisionOptimAMPConf(**self.amp)


Prompt:
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

# pyre-ignore-all-errors

import fnmatch
import itertools
import logging
from dataclasses import dataclass
from typing import Any, Dict, Iterable, List, Optional, Set, Tuple, Union

import hydra
import torch
import torch.nn as nn
from omegaconf import DictConfig, MISSING

from . import LARS, OmniOptimizer

def create_lars_optimizer(params, opt, **lars_params):
    optim = hydra.utils.instantiate(opt, params=params)
    return LARS(optim, **lars_params)

def validate_param_group_params(param_groups, model):

    parameters = [set(param_group["params"]) for param_group in param_groups]
    model_parameters = {parameter for _, parameter in model.named_parameters()}
    for p1, p2 in itertools.permutations(parameters, 2):
        assert p1.isdisjoint(p2), "Scheduler generated param_groups should be disjoint"
    assert (
        set.union(*parameters) == model_parameters
    ), "Scheduler generated param_groups include all parameters of the model"


Prompt:
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

# pyre-ignore-all-errors

import fnmatch
import itertools
import logging
from dataclasses import dataclass
from typing import Any, Dict, Iterable, List, Optional, Set, Tuple, Union

import hydra
import torch
import torch.nn as nn
from omegaconf import DictConfig, MISSING

from . import LARS, OmniOptimizer

def create_lars_optimizer(params, opt, **lars_params):
    optim = hydra.utils.instantiate(opt, params=params)
    return LARS(optim, **lars_params)

def validate_param_group_params(param_groups, model):
    parameters = [set(param_group["params"]) for param_group in param_groups]
    model_parameters = {parameter for _, parameter in model.named_parameters()}
    for p1, p2 in itertools.permutations(parameters, 2):
        assert p1.isdisjoint(p2), "Scheduler generated param_groups should be disjoint"
    assert (
        set.union(*parameters) == model_parameters
    ), "Scheduler generated param_groups include all parameters of the model"

def unix_pattern_to_parameter_names(
    scheduler_cfg: DictConfig, model: nn.Module
) -> Union[None, Set[str]]:

    if "param_names" not in scheduler_cfg and "module_cls_names" not in scheduler_cfg:
        return None
    return unix_param_pattern_to_parameter_names(scheduler_cfg, model).union(
        unix_module_cls_pattern_to_parameter_names(scheduler_cfg, model)
    )


Prompt:
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

# pyre-ignore-all-errors

import fnmatch
import itertools
import logging
from dataclasses import dataclass
from typing import Any, Dict, Iterable, List, Optional, Set, Tuple, Union

import hydra
import torch
import torch.nn as nn
from omegaconf import DictConfig, MISSING

from . import LARS, OmniOptimizer

def create_lars_optimizer(params, opt, **lars_params):
    optim = hydra.utils.instantiate(opt, params=params)
    return LARS(optim, **lars_params)

def validate_param_group_params(param_groups, model):
    parameters = [set(param_group["params"]) for param_group in param_groups]
    model_parameters = {parameter for _, parameter in model.named_parameters()}
    for p1, p2 in itertools.permutations(parameters, 2):
        assert p1.isdisjoint(p2), "Scheduler generated param_groups should be disjoint"
    assert (
        set.union(*parameters) == model_parameters
    ), "Scheduler generated param_groups include all parameters of the model"

def unix_pattern_to_parameter_names(
    scheduler_cfg: DictConfig, model: nn.Module
) -> Union[None, Set[str]]:
    if "param_names" not in scheduler_cfg and "module_cls_names" not in scheduler_cfg:
        return None
    return unix_param_pattern_to_parameter_names(scheduler_cfg, model).union(
        unix_module_cls_pattern_to_parameter_names(scheduler_cfg, model)
    )

def get_full_parameter_name(module_name, param_name):

    if module_name == "":
        return param_name
    return f"{module_name}.{param_name}"


Prompt:
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

# pyre-ignore-all-errors

import fnmatch
import itertools
import logging
from dataclasses import dataclass
from typing import Any, Dict, Iterable, List, Optional, Set, Tuple, Union

import hydra
import torch
import torch.nn as nn
from omegaconf import DictConfig, MISSING

from . import LARS, OmniOptimizer

def create_lars_optimizer(params, opt, **lars_params):
    optim = hydra.utils.instantiate(opt, params=params)
    return LARS(optim, **lars_params)

def validate_param_group_params(param_groups, model):
    parameters = [set(param_group["params"]) for param_group in param_groups]
    model_parameters = {parameter for _, parameter in model.named_parameters()}
    for p1, p2 in itertools.permutations(parameters, 2):
        assert p1.isdisjoint(p2), "Scheduler generated param_groups should be disjoint"
    assert (
        set.union(*parameters) == model_parameters
    ), "Scheduler generated param_groups include all parameters of the model"

def unix_pattern_to_parameter_names(
    scheduler_cfg: DictConfig, model: nn.Module
) -> Union[None, Set[str]]:
    if "param_names" not in scheduler_cfg and "module_cls_names" not in scheduler_cfg:
        return None
    return unix_param_pattern_to_parameter_names(scheduler_cfg, model).union(
        unix_module_cls_pattern_to_parameter_names(scheduler_cfg, model)
    )

def get_full_parameter_name(module_name, param_name):
    if module_name == "":
        return param_name
    return f"{module_name}.{param_name}"

def unix_module_cls_pattern_to_parameter_names(
    scheduler_cfg: DictConfig,
    model: nn.Module,
) -> Union[None, Set[str]]:
    if "module_cls_names" not in scheduler_cfg:
        return set()
    module_cls_to_params = {}
    for module_name, module in model.named_modules():
        module_cls = type(module)
        module_cls_to_params.setdefault(module_cls, set())
        module_cls_to_params[module_cls] |= set(
            get_full_parameter_name(module_name, param_name)
            for param_name, _ in module.named_parameters()
        )
    parameter_names = []
    for module_cls_name in scheduler_cfg.module_cls_names:
        module_cls = hydra.utils.get_class(module_cls_name)
        matching_parameters = module_cls_to_params.get(module_cls, set())
        assert len(matching_parameters) > 0, (
            f"Optimizer option for {scheduler_cfg.option} module_cls_name"
            f" {module_cls_name} does not match any classes in the model"
        )
        logging.info(
            f"Matches for module_cls_name [{module_cls_name}]: {matching_parameters} "
        )
        parameter_names.append(matching_parameters)
    return set.union(*parameter_names)

def unix_param_pattern_to_parameter_names(
    scheduler_cfg: DictConfig,
    model: nn.Module,
) -> Union[None, Set[str]]:

    if "param_names" not in scheduler_cfg:
        return set()
    all_parameter_names = {name for name, _ in model.named_parameters()}
    parameter_names = []
    for param_name in scheduler_cfg.param_names:
        matching_parameters = set(fnmatch.filter(all_parameter_names, param_name))
        assert len(matching_parameters) >= 1, (
            f"Optimizer option for {scheduler_cfg.option} param_names {param_name} "
            "does not match any parameters in the model"
        )
        logging.info(f"Matches for param_name [{param_name}]: {matching_parameters}")
        parameter_names.append(matching_parameters)
    return set.union(*parameter_names)


Prompt:
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

# pyre-ignore-all-errors

import fnmatch
import itertools
import logging
from dataclasses import dataclass
from typing import Any, Dict, Iterable, List, Optional, Set, Tuple, Union

import hydra
import torch
import torch.nn as nn
from omegaconf import DictConfig, MISSING

from . import LARS, OmniOptimizer

def create_lars_optimizer(params, opt, **lars_params):
    optim = hydra.utils.instantiate(opt, params=params)
    return LARS(optim, **lars_params)

def validate_param_group_params(param_groups, model):
    parameters = [set(param_group["params"]) for param_group in param_groups]
    model_parameters = {parameter for _, parameter in model.named_parameters()}
    for p1, p2 in itertools.permutations(parameters, 2):
        assert p1.isdisjoint(p2), "Scheduler generated param_groups should be disjoint"
    assert (
        set.union(*parameters) == model_parameters
    ), "Scheduler generated param_groups include all parameters of the model"

def unix_pattern_to_parameter_names(
    scheduler_cfg: DictConfig, model: nn.Module
) -> Union[None, Set[str]]:
    if "param_names" not in scheduler_cfg and "module_cls_names" not in scheduler_cfg:
        return None
    return unix_param_pattern_to_parameter_names(scheduler_cfg, model).union(
        unix_module_cls_pattern_to_parameter_names(scheduler_cfg, model)
    )

def get_full_parameter_name(module_name, param_name):
    if module_name == "":
        return param_name
    return f"{module_name}.{param_name}"

def unix_module_cls_pattern_to_parameter_names(
    scheduler_cfg: DictConfig,
    model: nn.Module,
) -> Union[None, Set[str]]:
    if "module_cls_names" not in scheduler_cfg:
        return set()
    module_cls_to_params = {}
    for module_name, module in model.named_modules():
        module_cls = type(module)
        module_cls_to_params.setdefault(module_cls, set())
        module_cls_to_params[module_cls] |= set(
            get_full_parameter_name(module_name, param_name)
            for param_name, _ in module.named_parameters()
        )
    parameter_names = []
    for module_cls_name in scheduler_cfg.module_cls_names:
        module_cls = hydra.utils.get_class(module_cls_name)
        matching_parameters = module_cls_to_params.get(module_cls, set())
        assert len(matching_parameters) > 0, (
            f"Optimizer option for {scheduler_cfg.option} module_cls_name"
            f" {module_cls_name} does not match any classes in the model"
        )
        logging.info(
            f"Matches for module_cls_name [{module_cls_name}]: {matching_parameters} "
        )
        parameter_names.append(matching_parameters)
    return set.union(*parameter_names)

def unix_param_pattern_to_parameter_names(
    scheduler_cfg: DictConfig,
    model: nn.Module,
) -> Union[None, Set[str]]:
    if "param_names" not in scheduler_cfg:
        return set()
    all_parameter_names = {name for name, _ in model.named_parameters()}
    parameter_names = []
    for param_name in scheduler_cfg.param_names:
        matching_parameters = set(fnmatch.filter(all_parameter_names, param_name))
        assert len(matching_parameters) >= 1, (
            f"Optimizer option for {scheduler_cfg.option} param_names {param_name} "
            "does not match any parameters in the model"
        )
        logging.info(f"Matches for param_name [{param_name}]: {matching_parameters}")
        parameter_names.append(matching_parameters)
    return set.union(*parameter_names)

def set_default_parameters(
    scheduler_cfgs: List[DictConfig], all_parameter_names: Set[str]
) -> None:

    constraints = [
        scheduler_cfg.parameter_names
        for scheduler_cfg in scheduler_cfgs
        if scheduler_cfg.parameter_names is not None
    ]
    if len(constraints) == 0:
        default_params = set(all_parameter_names)
    else:

        default_params = all_parameter_names - set.union(*constraints)
    default_count = 0
    for scheduler_cfg in scheduler_cfgs:
        if scheduler_cfg.parameter_names is None:
            scheduler_cfg.parameter_names = default_params
            default_count += 1
    assert default_count <= 1, "Only one scheduler per option can be default"
    if default_count == 0:  # Add defaults without options
        scheduler_cfgs.append({"parameter_names": default_params})


Prompt:
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

# pyre-ignore-all-errors

import fnmatch
import itertools
import logging
from dataclasses import dataclass
from typing import Any, Dict, Iterable, List, Optional, Set, Tuple, Union

import hydra
import torch
import torch.nn as nn
from omegaconf import DictConfig, MISSING

from . import LARS, OmniOptimizer

def create_lars_optimizer(params, opt, **lars_params):
    optim = hydra.utils.instantiate(opt, params=params)
    return LARS(optim, **lars_params)

def validate_param_group_params(param_groups, model):
    parameters = [set(param_group["params"]) for param_group in param_groups]
    model_parameters = {parameter for _, parameter in model.named_parameters()}
    for p1, p2 in itertools.permutations(parameters, 2):
        assert p1.isdisjoint(p2), "Scheduler generated param_groups should be disjoint"
    assert (
        set.union(*parameters) == model_parameters
    ), "Scheduler generated param_groups include all parameters of the model"

def unix_pattern_to_parameter_names(
    scheduler_cfg: DictConfig, model: nn.Module
) -> Union[None, Set[str]]:
    if "param_names" not in scheduler_cfg and "module_cls_names" not in scheduler_cfg:
        return None
    return unix_param_pattern_to_parameter_names(scheduler_cfg, model).union(
        unix_module_cls_pattern_to_parameter_names(scheduler_cfg, model)
    )

def get_full_parameter_name(module_name, param_name):
    if module_name == "":
        return param_name
    return f"{module_name}.{param_name}"

def unix_module_cls_pattern_to_parameter_names(
    scheduler_cfg: DictConfig,
    model: nn.Module,
) -> Union[None, Set[str]]:
    if "module_cls_names" not in scheduler_cfg:
        return set()
    module_cls_to_params = {}
    for module_name, module in model.named_modules():
        module_cls = type(module)
        module_cls_to_params.setdefault(module_cls, set())
        module_cls_to_params[module_cls] |= set(
            get_full_parameter_name(module_name, param_name)
            for param_name, _ in module.named_parameters()
        )
    parameter_names = []
    for module_cls_name in scheduler_cfg.module_cls_names:
        module_cls = hydra.utils.get_class(module_cls_name)
        matching_parameters = module_cls_to_params.get(module_cls, set())
        assert len(matching_parameters) > 0, (
            f"Optimizer option for {scheduler_cfg.option} module_cls_name"
            f" {module_cls_name} does not match any classes in the model"
        )
        logging.info(
            f"Matches for module_cls_name [{module_cls_name}]: {matching_parameters} "
        )
        parameter_names.append(matching_parameters)
    return set.union(*parameter_names)

def unix_param_pattern_to_parameter_names(
    scheduler_cfg: DictConfig,
    model: nn.Module,
) -> Union[None, Set[str]]:
    if "param_names" not in scheduler_cfg:
        return set()
    all_parameter_names = {name for name, _ in model.named_parameters()}
    parameter_names = []
    for param_name in scheduler_cfg.param_names:
        matching_parameters = set(fnmatch.filter(all_parameter_names, param_name))
        assert len(matching_parameters) >= 1, (
            f"Optimizer option for {scheduler_cfg.option} param_names {param_name} "
            "does not match any parameters in the model"
        )
        logging.info(f"Matches for param_name [{param_name}]: {matching_parameters}")
        parameter_names.append(matching_parameters)
    return set.union(*parameter_names)

def set_default_parameters(
    scheduler_cfgs: List[DictConfig], all_parameter_names: Set[str]
) -> None:
    constraints = [
        scheduler_cfg.parameter_names
        for scheduler_cfg in scheduler_cfgs
        if scheduler_cfg.parameter_names is not None
    ]
    if len(constraints) == 0:
        default_params = set(all_parameter_names)
    else:

        default_params = all_parameter_names - set.union(*constraints)
    default_count = 0
    for scheduler_cfg in scheduler_cfgs:
        if scheduler_cfg.parameter_names is None:
            scheduler_cfg.parameter_names = default_params
            default_count += 1
    assert default_count <= 1, "Only one scheduler per option can be default"
    if default_count == 0:  # Add defaults without options
        scheduler_cfgs.append({"parameter_names": default_params})

def name_constraints_to_parameters(
    param_constraints: List[Set[str]], model: torch.nn.Module
) -> List[torch.nn.Parameter]:
    matching_names = set.intersection(*param_constraints)
    return [value for name, value in model.named_parameters() if name in matching_names]

def map_scheduler_cfgs_to_param_groups(
    scheduler_cfgs_per_param_group: Iterable[List[Dict]], model: torch.nn.Module
) -> Tuple[List[Dict[Any, Any]], List[Dict[str, List[torch.nn.Parameter]]]]:

    schedulers = []
    param_groups = []
    for scheduler_cfgs in scheduler_cfgs_per_param_group:
        param_constraints = [
            scheduler_cfg["parameter_names"] for scheduler_cfg in scheduler_cfgs
        ]
        matching_parameters = name_constraints_to_parameters(param_constraints, model)
        if len(matching_parameters) == 0:  # If no overlap of parameters, skip
            continue
        schedulers_for_group = {
            scheduler_cfg["option"]: scheduler_cfg["scheduler"]
            for scheduler_cfg in scheduler_cfgs
            if "option" in scheduler_cfg
        }
        schedulers.append(schedulers_for_group)
        param_groups.append({"params": matching_parameters})
    return schedulers, param_groups


Prompt:
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

class OmniOptimizer(object):
    def __init__(self, optimizer, schedulers=None) -> None:

        self.optimizer = optimizer
        self.schedulers = schedulers
        self._validate_optimizer_schedulers()
        self.step_schedulers(0.0)


Prompt:
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

class OmniOptimizer(object):
    def __init__(self, optimizer, schedulers=None) -> None:
        self.optimizer = optimizer
        self.schedulers = schedulers
        self._validate_optimizer_schedulers()
        self.step_schedulers(0.0)

    def _validate_optimizer_schedulers(self):

        if self.schedulers is None:
            return
        for _, set_of_schedulers in enumerate(self.schedulers):
            for option, _ in set_of_schedulers.items():
                assert option in self.optimizer.defaults, (
                    "Optimizer option "
                    f"{option} not found in {self.optimizer}. Valid options are "
                    f"{self.optimizer.defaults.keys()}"
                )


Prompt:
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

class OmniOptimizer(object):
    def __init__(self, optimizer, schedulers=None) -> None:
        self.optimizer = optimizer
        self.schedulers = schedulers
        self._validate_optimizer_schedulers()
        self.step_schedulers(0.0)

    def _validate_optimizer_schedulers(self):
        if self.schedulers is None:
            return
        for _, set_of_schedulers in enumerate(self.schedulers):
            for option, _ in set_of_schedulers.items():
                assert option in self.optimizer.defaults, (
                    "Optimizer option "
                    f"{option} not found in {self.optimizer}. Valid options are "
                    f"{self.optimizer.defaults.keys()}"
                )

    def step_schedulers(self, where: float) -> None:

        if self.schedulers is None:
            return
        for i, param_group in enumerate(self.optimizer.param_groups):
            for option, scheduler in self.schedulers[i].items():
                new_value = scheduler(where)
                param_group[option] = new_value


Prompt:
from __future__ import annotations

from ..checks import _check_type
from ..types import Column, Func, PandasDataFrame

def mutate(df: PandasDataFrame, over: dict[Column, Func]) -> PandasDataFrame:

    _check_type(over, dict)
    df = df.copy()
    for column, mutation in over.items():
        df[column] = df.apply(mutation, axis=1)
    return df  # type: ignore


Prompt:
from __future__ import annotations

from ..checks import _check_type
from ..types import LazyColumns, PandasDataFrame

def drop(df: PandasDataFrame, columns: LazyColumns) -> PandasDataFrame:

    _check_type(columns, {list, str})
    df = df.drop(columns, axis=1)
    return df


Prompt:
from ..checks import _check_type, _check_values
from ..types import NewColumn, OldColumn, PandasDataFrame

def rename(df: PandasDataFrame, columns: dict[OldColumn, NewColumn]) -> PandasDataFrame:

    _check_type(columns, dict)
    cv = columns.values()
    _check_values(cv, str)
    if len(set(cv)) != len(cv):
        raise KeyError("columns must be unique")
    missing_keys = set(columns.keys()) - set(df.columns)
    if missing_keys and len(missing_keys) == 1:
        raise KeyError(f"column key ({missing_keys}) is invalid")
    if missing_keys and len(missing_keys) > 1:
        raise KeyError(f"column keys ({missing_keys}) are invalid")
    df = df.rename(columns=columns)
    return df


Prompt:
import uuid

from ..checks import _check_type
from ..types import Column, Columns, PandasDataFrame

def split(
    df: PandasDataFrame, column: Column, into: Columns, sep: str, drop: bool = True
) -> PandasDataFrame:

    _check_type(column, str)
    _check_type(into, list)
    _check_type(sep, str)
    _check_type(drop, bool)
    if len(into) != len(set(into)):
        raise KeyError("into keys must be unique")
    if (column in into) and (not drop):
        raise KeyError("into keys must be unique")
    bad_keys = set(df.columns).difference(set([column])).intersection(set(into))
    if bad_keys:
        raise KeyError("into keys must be unique")
    columns = {uuid.uuid4().hex: col for col in into}
    temp = list(columns.keys())
    df = df.copy()
    df[temp] = df[column].str.split(sep, expand=True)
    if drop:
        df = df.drop(column, axis=1)
    df = df.rename(columns=columns)
    return df


Prompt:
from __future__ import annotations

from ..checks import _check_keys, _check_type
from ..types import LazyColumns, PandasDataFrame

def sort(
    df: PandasDataFrame, columns: LazyColumns, descending: bool = False
) -> PandasDataFrame:

    _check_type(columns, {list, str})
    _check_type(descending, bool)
    _check_keys(columns, df.columns)
    df = df.sort_values(by=columns, ascending=not descending)
    df = df.reset_index(drop=True)
    return df


Prompt:
from __future__ import annotations

from ..checks import _check_type
from ..types import Column, PandasDataFrame, PandasGroupedFrame

def pack(
    df: PandasDataFrame | PandasGroupedFrame, column: Column, sep: str
) -> PandasDataFrame:

    _check_type(column, str)
    _check_type(sep, str)
    order = df.obj.columns if isinstance(df, PandasGroupedFrame) else df.columns  # type: ignore
    df = df.agg(**{column: (column, lambda x: x.astype(str).str.cat(sep=sep))})  # type: ignore
    df = df[[col for col in df.columns if col in order]]
    df = df.reset_index(drop=True)
    return df


Prompt:
from __future__ import annotations

from ..checks import _check_type
from ..types import Column, Func, PandasDataFrame, PandasGroupedFrame

def rollup(
    df: PandasDataFrame | PandasGroupedFrame,
    over: dict[Column, tuple[Column, Func]],
) -> PandasDataFrame:

    _check_type(over, dict)
    if isinstance(df, PandasGroupedFrame):
        groups = set(df.grouper.names)  # type: ignore
        keys = set(over.keys())
        if groups.intersection(keys):
            raise KeyError("unable to overwrite group keys")
        df = df.agg(**over)
        df = df.reset_index(drop=True)
    else:
        df = df.agg(**over)  # type: ignore
        df = df.T  # type: ignore
        df = df.reset_index(drop=True)  # type: ignore
        df = df.fillna(method="ffill")  # type: ignore
        df = df.fillna(method="bfill")  # type: ignore
        df = df.head(1)  # type: ignore
    return df


Prompt:
import pandas as pd  # pyright: ignore[reportMissingImports]

from ..types import PandasDataFrame

def append(top: PandasDataFrame, bottom: PandasDataFrame) -> PandasDataFrame:

    df = pd.concat([top, bottom])
    df = df.reset_index(drop=True)
    return df


Prompt:
from __future__ import annotations

import warnings

from ..checks import _check_type
from ..types import Column, Columns, PandasDataFrame

def combine(
    df: PandasDataFrame, columns: Columns, into: Column, sep: str, drop: bool = True
) -> PandasDataFrame:

    _check_type(columns, list)
    _check_type(into, str)
    _check_type(sep, str)
    _check_type(drop, bool)
    into_is_in_columns = into in columns
    into_is_not_in_columns = not into_is_in_columns
    into_is_in_df_columns = into in df.columns
    if into_is_not_in_columns and into_is_in_df_columns:
        message = f"overwriting existing column '{into}'"
        warnings.warn(message)
    df = df.copy()
    new = df[columns].apply(lambda row: sep.join(row.values.astype(str)), axis=1)
    if drop:
        df = df.drop(columns, axis=1)
    df[into] = new
    return df


Prompt:
from ..checks import _check_type
from ..types import Column, NewValue, OldValue, PandasDataFrame

def replace(
    df: PandasDataFrame, over: dict[Column, dict[OldValue, NewValue]]
) -> PandasDataFrame:

    _check_type(over, dict)
    bad_columns = list(set(over.keys()) - set(df.columns))
    if bad_columns and len(bad_columns) == 1:
        raise KeyError(f"column key: {bad_columns} is invalid")
    if bad_columns and len(bad_columns) > 1:
        raise KeyError(f"column keys: {bad_columns} are invalid")
    df = df.replace(over)
    return df


Prompt:
from __future__ import annotations

from ..checks import _check_type
from ..types import LazyColumns, PandasDataFrame

def denix(df: PandasDataFrame, columns: LazyColumns | None = None) -> PandasDataFrame:

    _check_type(columns, {list, str, None})
    columns = [columns] if isinstance(columns, str) else columns
    if isinstance(columns, list):
        bad_keys = set(columns).difference(df.columns)
        if bad_keys:
            if len(bad_keys) == 1:
                message = f"columns argument contains invalid key {bad_keys}"
            else:
                message = f"columns argument contains invalid keys {bad_keys}"
            raise KeyError(message)
    df = df.dropna(subset=columns)
    df = df.reset_index(drop=True)
    return df


Prompt:
from __future__ import annotations

from ..checks import _check_keys, _check_type
from ..types import LazyColumns, PandasDataFrame

def dedupe(df: PandasDataFrame, columns: LazyColumns | None = None) -> PandasDataFrame:

    _check_type(columns, {list, str, None})
    _check_keys(columns, df.columns)
    df = df.drop_duplicates(subset=columns, keep="first")
    df = df.reset_index(drop=True)
    return df


Prompt:
from __future__ import annotations

import warnings

from ..checks import _check_type
from ..types import Column, PandasDataFrame, PandasGroupedFrame

def accumulate(
    df: PandasDataFrame | PandasGroupedFrame, column: Column, into: Column
) -> PandasDataFrame:

    _check_type(column, str)
    _check_type(into, str)
    if isinstance(df, PandasDataFrame):
        into_is_not_column = into != column
        into_is_in_df_columns = into in df.columns
        if into_is_not_column and into_is_in_df_columns:
            message = f"overwriting existing column '{into}'"
            warnings.warn(message)
        df = df.copy()
    result = df[column].cumsum()
    if isinstance(df, PandasGroupedFrame):
        df = df.obj.copy()  # type: ignore
    df[into] = result  # type: ignore
    return df  # type: ignore


Prompt:
from __future__ import annotations

import pandas as pd  # pyright: ignore[reportMissingImports]

from ..checks import _check_type
from ..types import PandasDataFrame

def cross(
    lhs: PandasDataFrame,
    rhs: PandasDataFrame,
    postfix: tuple[str, str] = ("_lhs", "_rhs"),
) -> PandasDataFrame:

    _check_type(postfix, tuple)
    df = pd.merge(lhs, rhs, how="cross", suffixes=postfix)
    df = df.reset_index(drop=True)
    return df


Prompt:
from __future__ import annotations

from ..checks import _check_type
from ..types import PandasDataFrame

def sample(
    df: PandasDataFrame, rows: int | float, seed: int | None = None
) -> PandasDataFrame:

    _check_type(rows, {int, float})
    if rows >= 1:
        if isinstance(rows, float):
            raise ValueError("must be int if > 1")
        df = df.sample(rows, random_state=seed)
    elif 0 < rows < 1:
        df = df.sample(frac=rows, random_state=seed)
    else:
        raise ValueError("must be > 0")
    df = df.reset_index(drop=True)
    return df


Prompt:
from __future__ import annotations

from ..checks import _check_type
from ..types import PandasDataFrame

def shuffle(df: PandasDataFrame, seed: int | None = None) -> PandasDataFrame:

    _check_type(seed, {int, None})
    df = df.sample(frac=1, random_state=seed)
    df = df.reset_index(drop=True)
    return df


Prompt:
from __future__ import annotations

from ..checks import _check_type
from ..types import Column, PandasDataFrame

def unpack(df: PandasDataFrame, column: Column, sep: str) -> PandasDataFrame:

    _check_type(column, str)
    _check_type(sep, str)
    df = df.assign(**{column: df[column].str.split(sep)})
    df = df.explode(column)
    df = df.reset_index(drop=True)
    return df


Prompt:
from __future__ import annotations

import warnings

from ..checks import _check_type
from ..types import Column, PandasDataFrame, PandasGroupedFrame

def rank(
    df: PandasDataFrame | PandasGroupedFrame,
    column: Column,
    into: Column,
    descending: bool = False,
) -> PandasDataFrame:

    _check_type(column, str)
    _check_type(into, str)
    _check_type(descending, bool)
    if isinstance(df, PandasDataFrame):
        into_is_not_column = into != column
        into_is_in_df_columns = into in df.columns
        if into_is_not_column and into_is_in_df_columns:
            message = f"overwriting existing column '{into}'"
            warnings.warn(message)
        df = df.copy()
    result = df[column].rank(method="dense", ascending=not descending)
    if isinstance(df, PandasGroupedFrame):
        df = df.obj.copy()  # type: ignore
    df[into] = result  # type: ignore
    return df  # type: ignore


Prompt:
from __future__ import annotations

from ..checks import _check_type
from ..types import Direction, LazyColumns, PandasDataFrame, Value

def fill(
    df: PandasDataFrame,
    columns: LazyColumns | None = None,
    direction: Direction | None = None,
    constant: Value | None = None,
) -> PandasDataFrame:

    _check_type(columns, {list, str, None})
    _check_type(direction, {str, None})
    columns = [columns] if isinstance(columns, str) else columns
    if (direction != None) and (constant != None):
        raise ValueError("either direction OR constant must be None")
    if (direction == None) and (constant == None):
        raise ValueError("either direction OR constant must not be None")
    if direction != None:
        if not (direction in ["down", "up"]):
            raise ValueError("must be one of {'down', 'up'}")
        method = {"down": "ffill", "up": "bfill"}.get(direction)
        value = None
    if constant != None:
        value = constant
        method = None
    df = df.copy()
    if columns:
        df[columns] = df[columns].fillna(value=value, method=method)  # type: ignore
    else:
        df = df.fillna(value=value, method=method)  # type: ignore
    return df