[Question] Creating codebooks for Q, K, V vectors

[cmathw]

In codebook_features/models.py, I can see a method for attaching codebooks to each attention block's query, key and value vectors:

https://github.com/taufeeque9/codebook-features/blob/a37ea8fe7d4d39298aaea042a078d09401396edc/codebook_features/models.py#L1439C1-L1460C36

After training a model with these codebooks attached though, it does not seem possible to convert this model to a HookedTransformer model (doing so raises AttributeError: 'HookedTransformer' object has no attribute 'qkv_key'). What is the status of using qkv codebooks and converting to a HookedTransformer model currently? Happy to write a PR if this needs integrating with HookedTransformerCodebookModel class to work.