def precedes_final_layer(
    derived_scalar_location_within_layer: LocationWithinLayer | None,
    derived_scalar_layer_index: LayerIndex,
    final_residual_location_within_layer: LocationWithinLayer | None,
    final_residual_layer_index: LayerIndex,
) -> bool:
    """Returns True if the derived scalar at a given layer_index precedes the final residual stream derived scalar
    at a specified layer_index"""
    # return True if the derived scalar at a given layer_index precedes the final residual stream layer_index
    if derived_scalar_layer_index is None:
        return True  # activations with no layer_index are assumed to precede
        # all activations with layer_index; note that according to current conventions
        # this is true for all residual stream activations (not true e.g. for token logits)
    elif final_residual_layer_index is None:
        assert derived_scalar_layer_index is not None
        return False  # activations with layer_index precede activations with no layer_index
    elif derived_scalar_layer_index < final_residual_layer_index:
        return True
    elif derived_scalar_layer_index == final_residual_layer_index:
        if derived_scalar_location_within_layer is None:
            raise ValueError(
                "derived_scalar_location_within_layer must be provided in case of equal layer indices"
            )
        if final_residual_location_within_layer is None:
            raise ValueError(
                "final_residual_location_within_layer must be provided in case of equal layer indices"
            )
        if derived_scalar_location_within_layer < final_residual_location_within_layer:
            # location_within_layer inherits from int; therefore they are straightforwardly comparable
            return True
        else:
            return False
    else:
        assert derived_scalar_layer_index > final_residual_layer_index
        return False
Returns True if the derived scalar at a given layer_index precedes the final residual stream derived scalar
    at a specified layer_index
neuron_explainer/activations/derived_scalars/locations.py