def get_dimension_index_of_weight_location_type(
    weight_location_type: WeightLocationType, dimension: Dimension
) -> int:
    """Returns the index of a dimension within a weight tensor, and raises an error if the
    dimension is found 0 or >1 time (0 can happen, but >1 indicates a bug somewhere)."""
    assert weight_shape_by_location_type[weight_location_type].count(dimension) == 1
    return weight_shape_by_location_type[weight_location_type].index(dimension)
Returns the index of a dimension within a weight tensor, and raises an error if the
    dimension is found 0 or >1 time (0 can happen, but >1 indicates a bug somewhere).
neuron_explainer/models/model_component_registry.py