def location_within_layer()

in neuron_explainer/models/model_component_registry.py [0:0]

• 19 lines of code
• 6 McCabe index (conditional complexity)

    def location_within_layer(self) -> Optional["LocationWithinLayer"]:
        # this uses the information available to infer the location within a layer of a specific node_type.
        # It returns None in cases where the location within layer is ambiguous based on the information
        # provided; e.g. for residual stream node types, it might be post-attn or post-mlp.
        # (with activation location type
        # for additional clarification). It returns None in cases where the location within layer is ambiguous based on the information
        # provided; one example is for autoencoder latents, which might be based on any dst. In this case, further information from the
        # DSTConfig is used.
        # It throws an error if the node_type is *never* associated with a location within layer (e.g. vocab tokens)
        match self:
            case NodeType.MLP_NEURON:
                return LocationWithinLayer.MLP
            case NodeType.ATTENTION_HEAD | NodeType.QK_CHANNEL | NodeType.V_CHANNEL:
                return LocationWithinLayer.ATTN
            case (
                NodeType.RESIDUAL_STREAM_CHANNEL
                | NodeType.LAYER
                | NodeType.AUTOENCODER_LATENT
                | NodeType.MLP_AUTOENCODER_LATENT
                | NodeType.ATTENTION_AUTOENCODER_LATENT
                | NodeType.AUTOENCODER_LATENT_BY_TOKEN_PAIR
            ):
                # these node types are ambiguous based on the information provided
                return None
            case NodeType.VOCAB_TOKEN:
                # users should not be asking about the location within layer of vocab tokens; this indicates something's wrong
                raise ValueError("Vocab tokens don't have a location within layer")
            case _:
                raise NotImplementedError(f"Unknown node type {self=}")