def compute_measure(
    model: nn.Module,
    init_model: nn.Module,
    measure_func: Callable,
    operator: str,
    kwargs: dict = {},
    p: int = 1,
) -> float:
    """
    Computes measure value for each layer given trained network and network at
    initialization.  Then aggregates values per layer using specified operator.

    :param model: trained network
    :param init_model: network at initialization
    :param measure_func: callable for the measure to compute
    :param operator: 'log_product', 'sum', 'max', 'product', or 'norm'
    :param p: p in L^p
    :return: value of the desired measure
    """

    measure_value = 0
    # weight_modules = ["Linear", "Embedding"]
    weight_modules = ["Linear"]

    if operator == "product":
        measure_value = math.exp(
            compute_measure(model, init_model, measure_func, "log_product", kwargs, p)
        )
    elif operator == "norm":
        measure_value = (
            compute_measure(model, init_model, measure_func, "sum", kwargs, p=p)
        ) ** (1 / p)
    else:
        measure_value = 0
        for child, init_child in zip(model.children(), init_model.children()):
            module_name = child._get_name()
            if module_name in weight_modules:
                if operator == "log_product":
                    measure_value += math.log(measure_func(child, init_child, **kwargs))
                elif operator == "sum":
                    measure_value += (measure_func(child, init_child, **kwargs)) ** p
                elif operator == "max":
                    measure_value = max(
                        measure_value, measure_func(child, init_child, **kwargs)
                    )
            else:
                measure_value += compute_measure(
                    child, init_child, measure_func, operator, kwargs, p=p
                )
    return measure_value

    Computes measure value for each layer given trained network and network at
    initialization.  Then aggregates values per layer using specified operator.

    :param model: trained network
    :param init_model: network at initialization
    :param measure_func: callable for the measure to compute
    :param operator: 'log_product', 'sum', 'max', 'product', or 'norm'
    :param p: p in L^p
    :return: value of the desired measure
    
grok/metrics.py