pyvene.models.gpt_neo.modelings_intervenable_gpt_neo.gpt_neo_type_to_dimension_mapping

gpt_neo_type_to_dimension_mapping = {'attention_input': ('hidden_size',), 'attention_output': ('hidden_size',), 'attention_value_output': ('hidden_size',), 'block_input': ('hidden_size',), 'block_output': ('hidden_size',), 'head_attention_value_output': ('hidden_size/num_heads',), 'head_key_output': ('hidden_size/num_heads',), 'head_query_output': ('hidden_size/num_heads',), 'head_value_output': ('hidden_size/num_heads',), 'key_output': ('hidden_size',), 'mlp_activation': ('intermediate_size', 'hidden_size*4'), 'mlp_input': ('hidden_size',), 'mlp_output': ('hidden_size',), 'n_head': 'num_heads', 'query_output': ('hidden_size',), 'value_output': ('hidden_size',)}

gpt_neo model with LM head