from .attn_proj_block import ArtiHippoBlock
from .qkv_block import ArtiHippoQueryBlock, ArtiHippoKeyBlock, ArtiHippoValueBlock
from .mlp_block import ArtiHippoMlpBlock


ARTIHIPPO_BLOCKS = {
    "attn_proj": ArtiHippoBlock,
    "query": ArtiHippoQueryBlock,
    "key": ArtiHippoKeyBlock,
    "value": ArtiHippoValueBlock,
    "ffn": ArtiHippoMlpBlock
}
