pyvene.models.gru.modelings_intervenable_gru.gru_type_to_dimension_mapping

gru_type_to_dimension_mapping = {'cell_input': ('h_dim',), 'cell_output': ('h_dim',), 'h2h_output': ('h_dim*3',), 'new_gate_input': ('h_dim',), 'new_gate_output': ('h_dim',), 'new_h2h_output': ('h_dim',), 'new_x2h_output': ('h_dim',), 'reset_gate_input': ('h_dim',), 'reset_gate_output': ('h_dim',), 'reset_h2h_output': ('h_dim',), 'reset_x2h_output': ('h_dim',), 'update_gate_input': ('h_dim',), 'update_gate_output': ('h_dim',), 'update_h2h_output': ('h_dim',), 'update_x2h_output': ('h_dim',), 'x2h_output': ('h_dim*3',)}

mlp model with classification head