Make matrix.py capable handling weight-sharing dim

[wiseodd]

Currently KronLaplace doesn't work with arbitrary batch shape since matrix.py checks for at most 3-dim tensor. This is problematic when multiplying with Jacobians of shape (..., n_classes, n_params).

See https://github.com/aleximmer/Laplace/blob/glm-multidim/examples/lm_example.py for use case.

@runame, @aleximmer, any pointer on what to do? Is the fix simply to reshape into (-1, n_classes, n_params)?

Laplace/laplace/utils/matrix.py

Lines 208 to 215 in e0a68e5

	if W.ndim == 1:
	return self._bmm(W.unsqueeze(0).unsqueeze(0)).squeeze()
	elif W.ndim == 2:
	return self._bmm(W.unsqueeze(1)).squeeze()
	elif W.ndim == 3:
	return self._bmm(W)
	else:
	raise ValueError("Invalid shape for W")