def get_num_layer_for_convnext(var_name, variant='tiny'):
    """
    Divide [3, 3, 27, 3] layers into 12 groups; each group is three
    consecutive blocks, including possible neighboring downsample layers;
    adapted from https://github.com/microsoft/unilm/blob/master/beit/optim_factory.py
    """
    num_max_layer = 12
    if "stem" in var_name:
        return 0

    # note:  moved norm_layer outside of downsample module
    elif "downsample" in var_name or "norm_layer" in var_name:
        stage_id = int(var_name.split('.')[2])
        if stage_id == 0:
            layer_id = 0
        elif stage_id == 1 or stage_id == 2:
            layer_id = stage_id + 1
        elif stage_id == 3:
            layer_id = 12
        return layer_id

    elif "stages" in var_name:
        stage_id = int(var_name.split('.')[2])
        block_id = int(var_name.split('.')[4])
        if stage_id == 0 or stage_id == 1:
            layer_id = stage_id + 1
        elif stage_id == 2:
            if variant == 'tiny':
                layer_id = 3 + block_id
            else:
                layer_id = 3 + block_id // 3
        elif stage_id == 3:
            layer_id = 12
        return layer_id

    else:
        return num_max_layer + 1

    Divide [3, 3, 27, 3] layers into 12 groups; each group is three
    consecutive blocks, including possible neighboring downsample layers;
    adapted from https://github.com/microsoft/unilm/blob/master/beit/optim_factory.py
    
analyze/convnexts4nd/src/models/baselines/convnext_timm.py