def location_within_layer()

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

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

    def location_within_layer(self) -> LocationWithinLayer | None:
        # this uses the information available to infer the location within a layer of a specific node_type (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 needed.
        # It throws an error if the node_type is not associated with a location within layer.
        if self.node_type.location_within_layer is None:
            if self.node_type == NodeType.RESIDUAL_STREAM_CHANNEL:
                if self == ActivationLocationType.RESID_POST_EMBEDDING:
                    return None
                elif self == ActivationLocationType.RESID_DELTA_ATTN:
                    return LocationWithinLayer.ATTN
                elif self == ActivationLocationType.RESID_POST_ATTN:
                    return LocationWithinLayer.RESID_POST_ATTN
                elif self == ActivationLocationType.RESID_DELTA_MLP:
                    return LocationWithinLayer.MLP
                elif self == ActivationLocationType.RESID_POST_MLP:
                    return LocationWithinLayer.RESID_POST_MLP
                else:
                    return None
            else:
                return None
        else:
            return self.node_type.location_within_layer