Function reduce_func

scripts/analysis_utils.py:44–53

Return the corresponding reduction function.

(method)

    return mask

def reduce_func(method):
    """Return the corresponding reduction function."""
    if method == 'sum':
        return torch.sum
    elif method == 'mean':
        return torch.mean
    elif method == 'count':
        return lambda x: sum(x.size())
    else:
        raise NotImplementedError()

def attr_grad(tensor, reduce='sum'):
    """Calculate attribute gradient."""