Support for multidim outputs

[wiseodd]

Closes #252

Thanks to Kazuki for updating asdfghjkl on PyPI, we can now handle 3D tensor outputs (e.g. in LLMs) with GLM predictive.

[wiseodd]

WIP adding an article based on lm_example.py to the docs.

[wiseodd]

@runame @aleximmer ready to review!

[aleximmer]

Thanks for the useful addition. To make sure this breaks nothing, I would recommend the test mentioned in the review, otherwise LGTM.

[aleximmer]

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.

[wiseodd]

@runame @aleximmer do you know what happened here?

...
Ji, f = batch_gradient(self.model, closure, return_outputs=True)
print("Jac")
print(Ji)
...

with these setup:

X, Y = torch.randn(10, 3), torch.randn(10, 1)
X_multidim, Y_multidim = X.reshape(5, 2, 3), Y.reshape(5, 2, 1)

model_std = nn.Linear(3, 1)
model_multidim = deepcopy(model_std)

...

la = Laplace(model_multidim, ...)
la.fit(...)

la(X_multidim)

prints

Jac
tensor([[ 0.2302,  0.1006,  0.4197,  1.0000],
        [ 0.0000,  0.0000,  0.0000,  0.0000],
        [ 0.3055,  1.1884, -0.4010,  1.0000],
        [ 0.0000,  0.0000,  0.0000,  0.0000],
        [-2.0167,  1.4216, -0.5488,  1.0000],
        [ 0.0000,  0.0000,  0.0000,  0.0000],
        [-0.4076,  0.2810, -1.7831,  1.0000],
        [ 0.0000,  0.0000,  0.0000,  0.0000],
        [ 0.1297,  0.2202, -0.5024,  1.0000],
        [ 0.0000,  0.0000,  0.0000,  0.0000]]

[wiseodd]

Notice that the returned Ji has row [0 ... 0]. Seems like ASDL's issue?

[wiseodd]

To reproduce: uv run pytest -k "test_predictive_multidim