Support for multidim outputs #266
Conversation
|
WIP adding an article based on |
|
@runame @aleximmer ready to review! |
| }, | ||
| batch_size=[100], | ||
| ) # simulates a dataset | ||
| dl = DataLoader(ds, batch_size=BATCH_SIZE, shuffle=False, collate_fn=lambda x: x) |
There was a problem hiding this comment.
Does it not work with the default collate_fn?
|
|
||
| def _check_jacobians(self, Js: torch.Tensor) -> None: | ||
| if not isinstance(Js, torch.Tensor): | ||
| raise ValueError("Jacobians have to be torch.Tensor.") | ||
| if not Js.device == self._device: | ||
| raise ValueError("Jacobians need to be on the same device as Laplace.") | ||
| m, k, p = Js.size() | ||
| if p != self.n_params: | ||
| if Js.shape[-1] != self.n_params: |
There was a problem hiding this comment.
For consistency with previous assertion, perhaps add and Js.ndim >= 3?
| @@ -1335,7 +1397,7 @@ def _nn_predictive_classification( | |||
|
|
|||
| vector_to_parameters(self.mean, self.params) | |||
|
|
|||
| return py | |||
| return torch.as_tensor(py) | |||
There was a problem hiding this comment.
Just to make sure: I suppose this is for typing?
| return_outputs=True, | ||
| batch_size=self._get_batch_size(x), | ||
| ) | ||
| Ji, f = batch_gradient(self.model, closure, return_outputs=True) |
There was a problem hiding this comment.
Why is the batch_size argument no needed anymore, or perhaps, why was it even needed before?
There was a problem hiding this comment.
Could you add a test where you test for equivalence between a "trivial" multidimensional model and a previously supported model? This can be just a linear model that mapped previously from [num_data, D] -> [num_data, K] and now from [num_data / L, L, D] -> [num_data / L, L, K], that is, the data set is simply reshaped. This should give an equivalent posterior and predictive. Maybe to make sure also have the option to use a NN for the test, not only the linear model case.
Closes #252
Thanks to Kazuki for updating
asdfghjklon PyPI, we can now handle 3D tensor outputs (e.g. in LLMs) with GLM predictive.