linearized_forward_diag API

[SamDuffield]

Currently we have an API like

vals, chol, aux = linearized_forward_diag(f, params, batch, sd_diag)

but perhaps and API like

vals, chol, aux = linearized_forward_diag(f, sd_diag)(params, batch)

might be cleaner as it provides a new function that retains the required signature of f. It's also better fitting with the torch.func API

[SamDuffield]

Additionally considerations:

• Should we add an additional observation_noise_sd (could be a different name) that corresponds to $\sigma$ in the predictive distribution $N( f \mid f^*, J^T \Sigma J + \sigma^2 \mathbb{I})$, details in https://arxiv.org/abs/1906.11537
• Should we remove the internal no_grad to allow user more flexibility? Although majority of cases the function would still need to be wrapped in no_grad to avoid memory usage from undesired gradients