Project: 040_MOS_Model_Synergy_for_Test-Time_Adaptation_on_LiDA | Function: mos (iclr)

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:** Main MOS (Model Synergy) framework that implements test-time adaptation for 3D object detection by computing synergy weights using inverse Gram matrix and aggregating historical model checkpoints

**KEY TERMS:** {
  "G": "Gram matrix",
  "c": "synergy weights (w in paper)",
  "model_bank": "model bank F",
  "super_model": "super model f*_t"
}

**CONTEXT CODE:**     # Mock objects for testing
    class MockCfg:
        class MODEL:
            class POST_PROCESSING:
                RECALL_THRESH_LIST = [0.3, 0.5, 0.7]
        class SELF_TRAIN:
            get = lambda self, key, default=None: default
            SCORE_THRESH = [0.5, 0.5, 0.5]
            TAR = {'LOSS_WEIGHT': 1.0}
            get = lambda self, key, default=None: default
        class TTA:
            METHOD = 'mos'
            SAVE_CKPT = [100, 200, 300]
            SAVE_CKPT_INTERVAL = 50
            INSTANCE_EVAL = False
            get = lambda self, key, default=None: default
            class MOS_SETTING:
                AGGREGATE_START_CKPT = 100
                AGGREGATE_NUM = 5
                AGGREGATE_BANK = True
                BANK_START_CKPT = 100
        CLASS_NAMES = ['Car', 'Pedestrian', 'Cyclist']

    cfg = MockCfg()

    # Mock functions for testing
    def load_data_to_gpu(batch): pass
    def build_network(model_cfg, num_class, dataset):
        class MockModel:
            def load_params_from_file(self, filename, logger, report_logger=False, to_cpu=None): pass
            def cuda(self): return self
            def eval(self): return self
            def __call__(self, batch): return [], {}
            def named_parameters(self): return []
            def named_buffers(self): return []
        return MockModel()

    dist = None
    wandb = type('MockWandb', (), {'log': lambda x, **kwargs: None})()
    tqdm = type('MockTqdm', (), {'tqdm': lambda *args, **kwargs: iter(range(10))})()

    def clip_grad_norm_(*args, **kwargs): pass
    def save_checkpoint(*args, **kwargs): pass
    def checkpoint_state(*args, **kwargs): return {}
    def update_ema_variables(*args, **kwargs): pass
    def TTA_augmentation(dataset, batch, strength=None): return batch

PSEUDO_LABELS = {}
NEW_PSEUDO_LABELS = {}

def mos(model, val_loader, rank, leave_pbar, ps_label_dir, cur_epoch, model_func=None, lr_scheduler=None,
                       accumulated_iter=None, optim_cfg=None, tbar=None, total_it_each_epoch=None,
                       dataloader_iter=None, tb_log=None,ema_model=None, optimizer=None, ckpt_save_interval_iter=64, ckpt_save_dir=None,
                        logger=None,model_copy=None):
    raise NotImplementedError("Function implementation has been masked for testing")

def aggregate_model(model_path_list, model_weights, dataset, logger, dist, main_model):
    # agg_model = build_network(model_cfg=cfg.MODEL,num_class=len(cfg.CLASS_NAMES),dataset=dataset)
    # agg_model.cuda()
    # Clear all parameters
    # for agg_param in agg_model.parameters():
    #     agg_param.data.mul_(0)
    # # Clear BN
    # for agg_bf in agg_model.named_buffers():
    #     name, value = agg_bf
    #     if 'running_mean' in name or 'running_var' in name:
    #         value.data.mul_(0)
    raise NotImplementedError("Function implementation has been masked for testing")

def save_pseudo_label_batch(input_dict,
                            pred_dicts=None,
                            need_update=False):
    """
    Save pseudo label for give batch.
    If model is given, use model to inference pred_dicts,
    otherwise, directly use given pred_dicts.

    Args:
        input_dict: batch data read from dataloader
        pred_dicts: Dict if not given model.
            predict results to be generated pseudo label and saved
        need_update: Bool.
            If set to true, use consistency matching to update pseudo label
    """
    pos_ps_nmeter = common_utils.NAverageMeter(len(cfg.CLASS_NAMES))
    ign_ps_nmeter = common_utils.NAverageMeter(len(cfg.CLASS_NAMES))

    batch_size = len(pred_dicts)
    for b_idx in range(batch_size):
        pred_cls_scores = pred_iou_scores = None
        if 'pred_boxes' in pred_dicts[b_idx]:
            # Exist predicted boxes passing self-training score threshold
            pred_boxes = pred_dicts[b_idx]['pred_boxes'].detach().cpu().numpy()
            pred_labels = pred_dicts[b_idx]['pred_labels'].detach().cpu().numpy()
            pred_scores = pred_dicts[b_idx]['pred_scores'].detach().cpu().numpy()
            if 'pred_cls_scores' in pred_dicts[b_idx]:
                pred_cls_scores = pred_dicts[b_idx]['pred_cls_scores'].detach().cpu().numpy()
            if 'pred_iou_scores' in pred_dicts[b_idx]:
                pred_iou_scores = pred_dicts[b_idx]['pred_iou_scores'].detach().cpu().numpy()

            # remove boxes under negative threshold
            if cfg.SELF_TRAIN.get('NEG_THRESH', None):
                labels_remove_scores = np.array(cfg.SELF_TRAIN.NEG_THRESH)[pred_labels - 1]
                remain_mask = pred_scores >= labels_remove_scores
                pred_labels = pred_labels[remain_mask]
                pred_scores = pred_scores[remain_mask]
                pred_boxes = pred_boxes[remain_mask]
                if 'pred_cls_scores' in pred_dicts[b_idx]:
                    pred_cls_scores = pred_cls_scores[remain_mask]
                if 'pred_iou_scores' in pred_dicts[b_idx]:
                    pred_iou_scores = pred_iou_scores[remain_mask]

            labels_ignore_scores = np.array(cfg.SELF_TRAIN.SCORE_THRESH)[pred_labels - 1]
            ignore_mask = pred_scores < labels_ignore_scores
            pred_labels[ignore_mask] = -pred_labels[ignore_mask]

            if cfg.SELF_TRAIN.get('FIX_POS_NUM', None):
                expected_pos_num = pred_labels.shape[0] if pred_labels.shape[0] < cfg.SELF_TRAIN.FIX_POS_NUM else cfg.SELF_TRAIN.FIX_POS_NUM
                pred_labels[expected_pos_num:][pred_labels[expected_pos_num:] > 0] = - \
                pred_labels[expected_pos_num:][pred_labels[expected_pos_num:] > 0]

            gt_box = np.concatenate((pred_boxes,
                                     pred_labels.reshape(-1, 1),
                                     pred_scores.reshape(-1, 1)), axis=1)

        else:
            # no predicted boxes passes self-training score threshold
            gt_box = np.zeros((0, 9), dtype=np.float32)

        gt_infos = {
            'gt_boxes': gt_box,
            'cls_scores': pred_cls_scores,
            'iou_scores': pred_iou_scores,
            'memory_counter': np.zeros(gt_box.shape[0])
        }

        # record pseudo label to pseudo label dict
        if need_update:
            ensemble_func = getattr(memory_ensemble_utils, cfg.SELF_TRAIN.MEMORY_ENSEMBLE.NAME)
            gt_infos = memory_ensemble_utils.memory_ensemble(
                PSEUDO_LABELS[input_dict['frame_id'][b_idx]], gt_infos,
                cfg.SELF_TRAIN.MEMORY_ENSEMBLE, ensemble_func
            )

        # counter the number of ignore boxes for each class
        for i in range(ign_ps_nmeter.n):
            num_total_boxes = (np.abs(gt_infos['gt_boxes'][:, 7]) == (i+1)).sum()
            ign_ps_nmeter.update((gt_infos['gt_boxes'][:, 7] == -(i+1)).sum(), index=i)
            pos_ps_nmeter.update(num_total_boxes - ign_ps_nmeter.meters[i].val, index=i)



Complete the function implementation replacing "raise NotImplementedError".