pyvene.models.modeling_utils.scatter_neurons

pyvene.models.modeling_utils.scatter_neurons#

scatter_neurons(tensor_input, replacing_tensor_input, component, unit, unit_locations_as_list, model_type, model_config, use_fast, device=None)[source]#

Replace selected neurons in tensor_input by replacing_tensor_input.

Parameters:

tensor_input – tensors of shape (batch_size, sequence_length, …) if

unit is “pos” or “h”, tensors of shape (batch_size, num_heads, sequence_length, …) if unit is “h.pos” :param replacing_tensor_input: tensors of shape (batch_size, sequence_length, …) if unit is “pos” or “h”, tensors of shape (batch_size, num_heads, sequence_length, …) if unit is “h.pos”. :param component: types of intervention representations, such as “block_output” and “query_output” :param unit: the intervention units to gather. Units could be “h” - head number, “pos” - position in the sequence, or “dim” - a particular dimension in the embedding space. If intervening multiple units, they are ordered and separated by .. Currently only support “pos”, “h”, and “h.pos” units. :param unit_locations_as_list: tuple of lists of lists of positions to gather in tensor_input, according to the unit. :param model_type: Hugging Face Model Type :param model_config: Hugging Face Model Config :param use_fast: whether to use fast path (TODO: fast path condition) :return the in-place modified tensor_input