def test_inv_remap_state_dict(model_name: str):
    """
    Verify that we can convert a HF BigCode model to flash_attn and back.
    """

    state_dict = state_dict_from_pretrained(model_name)
    config = GPTBigCodeConfig.from_pretrained(model_name)

    flash_state_dict = remap_state_dict_hf_bigcode(state_dict, config)
    recovered_state_dict = inv_remap_state_dict_hf_bigcode(flash_state_dict, config)

    assert set(state_dict.keys()) == set(recovered_state_dict.keys())

    for k in state_dict.keys():
        assert state_dict[k].shape == recovered_state_dict[k].shape
        torch.testing.assert_close(state_dict[k], recovered_state_dict[k], rtol=1e-6, atol=1e-6)

    Verify that we can convert a HF BigCode model to flash_attn and back.
    
tests/models/test_bigcode.py