import torch

class CombinedOptimizer:
    """
    A wrapper class for managing multiple optimizers.
    This allows different parameter groups to be optimized by different algorithms.
    """
    def __init__(self, optimizers):
        self.optimizers = optimizers
        # Expose param_groups for learning rate schedulers
        self.param_groups = []
        for opt in self.optimizers:
            self.param_groups.extend(opt.param_groups)

    def zero_grad(self, set_to_none=True):
        """Calls zero_grad() for all optimizers."""
        for opt in self.optimizers:
            opt.zero_grad(set_to_none=set_to_none)

    def step(self):
        """Calls step() for all optimizers."""
        for opt in self.optimizers:
            opt.step()

    def state_dict(self):
        """Returns a list of state_dicts from all optimizers."""
        return [opt.state_dict() for opt in self.optimizers]

    def load_state_dict(self, state_dicts):
        """Loads state_dicts into each respective optimizer."""
        for opt, state in zip(self.optimizers, state_dicts):
            opt.load_state_dict(state)