Project: s41467-025-58055-3 | Function: load_dbb_graph (python)

system: You are an expert Python programmer implementing algorithms from scientific papers.
Complete the masked function based on the algorithm description.
Follow Python best practices and maintain consistent code style.
Return **only** the complete Python function code.

user: Complete this Python function:

**DESCRIPTION:** Loads and configures dynamic branch blocks for adaptive model architecture scaling in federated learning. Implements the resource-adaptive framework by dynamically adjusting model complexity based on computational constraints, allowing heterogeneous clients with varying resources to participate fairly.

**KEY TERMS:** {
  "dbb_list": "dynamic branch configuration",
  "DiverseBranchBlock": "adaptive architecture component",
  "branches": "computational complexity scaling",
  "conv_m": "convolutional layer adaptation"
}

**CONTEXT CODE:**         # remove the states of removed paramters
        assert len(param_groups) == len(self.optimizer.param_groups)
        for param_group, param_group_old in zip(param_groups, self.optimizer.param_groups):
            params, params_old = param_group['params'], param_group_old['params']
            params = set(params)
            for param_old in params_old:
                if param_old not in params:
                    if param_old in self.optimizer.state:
                        del self.optimizer.state[param_old]
            param_group_old['params'] = param_group['params']

    def adjust_model(self, device_capacity):
        records = {}
        for key in self._metric_records:
            records[key] = self._metric_records[key].avg
        # print('records:', records)
        
        metric_records_sorted = sorted(records.items(),
                                       key=lambda item: item[1],
                                       reverse=True)
        logger.info('metric: {}'.format(metric_records_sorted))
        self._grow(metric_records_sorted, device_capacity)

        self._reset_optimizer()

    def state_dict(self):
        # save dbb graph
        res = {}
        res['dbb_graph'] = self.dbb_graph()
        return res

    def load_state_dict(self, state_dict):
        if 'dbb_graph' in state_dict:
            self.load_dbb_graph(state_dict['dbb_graph'])


    def dbb_graph(self):
        dbb_list = []

        def traverse(parent, prefix=''):
            for k, m in parent.named_children():
                path = prefix + '.' + k if prefix != '' else k
                if isinstance(m, DiverseBranchBlock):
                    dbb_list.append((path, m.branches))
                traverse(m, prefix=path)

        traverse(self.model)
        print(dbb_list)
        return dbb_list

    def load_dbb_graph(self, dbb_list: list):
        """
        Loads and configures dynamic branch blocks for adaptive model architecture scaling in federated learning.

        This function implements the resource-adaptive framework by dynamically adjusting model complexity
        based on computational constraints, allowing heterogeneous clients with varying resources to
        participate fairly in federated healthcare learning scenarios.

        The method reconstructs dynamic diverse branch block (DBB) structures from a saved configuration,
        enabling personalized model architectures that adapt to individual client computational capacities
        while maintaining collaborative learning effectiveness.

        Args:
            dbb_list (list): List of tuples containing (module_path, branch_configuration)
                           defining the dynamic architecture structure

        Key algorithmic elements:
        - Dynamic architecture reconstruction
        - Conditional logic for branch configuration
        - Model complexity adaptation
        - Parameter management across heterogeneous nodes
        - Resource-constraint-aware model scaling

        Paper reference: "Our approach enhances model accuracy, safeguards patient privacy,
                        and promotes equitable access to trustworthy and efficient AI-driven healthcare solutions"
        Formula: Dynamic model architecture adaptation based on resource constraints

        Implementation details:
        - Validates model cleanliness (no existing DBB structures)
        - Iteratively reconstructs DBB blocks from configuration
        - Maintains original convolution layer properties
        - Resets optimizer to accommodate new parameters
        - Enables CUDA acceleration for computational efficiency
        """
        # MASKED: Dynamic branch block graph loading for resource-adaptive federated learning
        raise NotImplementedError("This function is masked for testing purposes. The function loads and configures dynamic branch blocks for adaptive model architecture scaling in federated learning, implementing the resource-adaptive framework by dynamically adjusting model complexity based on computational constraints, allowing heterogeneous clients with varying resources to participate fairly.")


def get_params(model, lr, weight_decay=1e-5, filter_bias_and_bn=True, sort_params=False):
    if weight_decay != 0 and filter_bias_and_bn:
        if hasattr(model, 'no_weight_decay'):
            skip_list = model.no_weight_decay()
            print(f'no weight decay: {skip_list}')
        else:
            skip_list = ()
        params = _add_weight_decay(model, lr, weight_decay, skip_list=skip_list, sort_params=sort_params)
        weight_decay = 0
    else:
        named_params = list(model.named_parameters())
        if sort_params:
            named_params.sort(key=lambda x: x[0])
        params = [x[1] for x in named_params]
        params = [{'params': params, 'initial_lr': lr}]
    return params


def _add_weight_decay(model, lr, weight_decay=1e-5, skip_list=(), sort_params=False):
    decay = []
    no_decay = []
    named_params = list(model.named_parameters())
    if sort_params:
        named_params.sort(key=lambda x: x[0])
    for name, param in named_params:
        if not param.requires_grad:
            continue  # frozen weights
        skip = False
        for skip_name in skip_list:
            if skip_name.startswith('[g]'):
                if skip_name[3:] in name:
                    skip = True
            elif name == skip_name:
                skip = True
        if len(param.shape) == 1 or name.endswith(".bias") or skip:
            no_decay.append(param)
        else:
            decay.append(param)
    return [
        {'params': no_decay, 'weight_decay': 0., 'initial_lr': lr},
        {'params': decay, 'weight_decay': weight_decay, 'initial_lr': lr}]


def build_optimizer(opt, model, lr, eps=1e-10, momentum=0.9, weight_decay=1e-5, filter_bias_and_bn=True, nesterov=True, sort_params=False):
    # params in dyrep must be sorted to make sure optimizer can correctly
    # load the states in resuming
    params = get_params(model, lr, weight_decay, filter_bias_and_bn, sort_params=sort_params)

    if opt == 'rmsprop':
        optimizer = optim.RMSprop(params, lr, eps=eps, weight_decay=weight_decay, momentum=momentum)
    elif opt == 'rmsproptf':
        optimizer = RMSpropTF(params, lr, eps=eps, weight_decay=weight_decay, momentum=momentum)
    elif opt == 'sgd':
        optimizer = optim.SGD(params, lr, momentum=momentum, weight_decay=weight_decay, nesterov=nesterov)
    elif opt == 'adamw':
        optimizer = optim.AdamW(params, lr, eps=eps, weight_decay=weight_decay)
    else:
        raise NotImplementedError(f'Optimizer {opt} not implemented.')
    return optimizer

Complete the function implementation replacing "raise NotImplementedError".