Decouple predictive from the Laplace class?

[wiseodd]

Useful for: Users who want to implement custom predictive approximations.

Issue: Currently, the predictive approximation is tightly coupled with the Laplace class. So, if the user wanted to implement a new predictive approximation, they have to dig deep into this class, and it might break something not to mention that it can be confusing.

Proposal:

• 2-steps predictive interface (function output and link predictives)

class FunctionPredictive:

    def __init__(self, ...):
        ...

    def __call__(self, x):
        ''' Return 2 arrays for means and vars '''
        raise NotImplementedError()


class LinearizedPredictive(FunctionPredictive):

    def __init__(self, laplace_net, ...):
        self.laplace_net = laplace_net
        ...

    def __call__(self, x):
        J = compute_jacobian(laplace_net, x)
        return laplace_net.map_prediction(x), J.T @ laplace_net.covmat @ J


class LinkPredictive:

    def __init__(self, ...):
        ...

    def __call__(self, f_mean, f_var):
        raise NotImplementedError()


class ProbitPredictive(LinkPredictive):

    def __init__(self, ...):
        ...

    def __call__(self, f_mean, f_var):
        return torch.sigmoid(f_mean / torch.sqrt(1 + pi/8 * f_var))

• Usage

linearized_pred = LinearizedPred()
probit_pred = ProbitPred()  # Set it to `None` if one does regression
laplace_net = Laplace(..., function_predictive=linearized_pred, link_predictive=probit_pred)
laplace_net.fit(train_loader)
laplace_net(x)  # Prediction using the specified predictives