def convnext_timm_tiny_s4nd_2d_to_3d(model, state_dict, ignore_head=True, jank=False):
    """
    inputs:
        model: nn.Module, the from 'scratch' model
        state_dict: dict, from the pretrained weights
        ignore_head: bool, whether to inflate weights in the head (or keep scratch weights).
            If number of classes changes (eg, imagenet to hmdb51), then you need to use this.

    return:
        state_dict: dict, update with inflated weights
    """

    # model_scratch_params_dict = dict(model.named_parameters())
    model_scratch_params_dict = {**dict(model.named_parameters()), **dict(model.named_buffers())}
    prefix = list(state_dict.keys())[0].split('.')[0]  # grab prefix in the keys for state_dict params

    new_state_dict = copy.deepcopy(state_dict)

    # for key in state_dict.keys():
    #     print(key)

    # breakpoint()


    for key in sorted(model_scratch_params_dict.keys()):


        # need to add the predix 'model' in convnext
        key_with_prefix = prefix + '.' + key

        # HACK
        old_key_with_prefix = key_with_prefix.replace("inv_w_real", "log_w_real")

        # print(key)
        # if '.kernel.L' in key:
        #     print(key, state_dict[old_key_with_prefix])


        if '.kernel.0' in key:
            # temporal dim is loaded from scratch
            print("found .kernel.0:", key)
            new_state_dict[key_with_prefix] = model_scratch_params_dict[key]

        elif '.kernel.1' in key:
            # This is the 1st kernel --> 0th kernel from pretrained model
            print("FOUND .kernel.1, putting kernel 0 into kernel 1", key)
            new_state_dict[key_with_prefix] = state_dict[old_key_with_prefix.replace(".kernel.1", ".kernel.0")]
        elif '.kernel.2' in key:
            print("FOUND .kernel.2, putting kernel 1 into kernel 2", key)
            new_state_dict[key_with_prefix] = state_dict[old_key_with_prefix.replace(".kernel.2", ".kernel.1")]
        elif ignore_head and 'head' in key:
            # ignore head weights
            print("found head key / parameter, ignore", key)
            new_state_dict[key_with_prefix] = model_scratch_params_dict[key]
        # keys match
        else:
            # check if mismatched shape, if so, need to inflate
            # this covers cases where we did not use s4 (eg, optionally use conv2d in downsample or the stem)
            try:
                if model_scratch_params_dict[key].ndim != state_dict[old_key_with_prefix].ndim:
                    print("matching keys, but shapes mismatched!  Need to inflate!", key)
                    # need the index [-3], 3rd from last, the temporal dim
                    index = -3
                    dim_len = model_scratch_params_dict[key].shape[index]
                    # loaded_params = state_dict[key_with_prefix].unsqueeze(index)  # unsqueeze
                    # new_state_dict[key_with_prefix] = torch.cat(dim_len * [loaded_params], axis=index)  # stack at this dim
                    new_state_dict[key_with_prefix] = repeat(state_dict[old_key_with_prefix], '... h w -> ... t h w', t=dim_len) # torch.cat(dim_len * [loaded_params], axis=index)  # stack at this dim
                else:
                    # matching case, shapes, match, load into new_state_dict as is
                    new_state_dict[key_with_prefix] = state_dict[old_key_with_prefix]
            # something went wrong, the keys don't actually match (and they should)!
            except:
                print("unmatched key", key)
                breakpoint()
                # continue

    return new_state_dict

    inputs:
        model: nn.Module, the from 'scratch' model
        state_dict: dict, from the pretrained weights
        ignore_head: bool, whether to inflate weights in the head (or keep scratch weights).
            If number of classes changes (eg, imagenet to hmdb51), then you need to use this.

    return:
        state_dict: dict, update with inflated weights
    
analyze/convnexts4nd/src/models/baselines/convnext_timm.py