Handle logits with more than 2 axes in torch.func diag and full backends

[wiseodd]

Currently, we still assume that logits.shape == (batch_size, n_classes). E.g.:

Laplace/laplace/curvature/curvature.py

Lines 332 to 334 in a4d3ed6

	ps = torch.softmax(f, dim=-1)
	G = torch.diag_embed(ps) - torch.einsum('mk,mc->mck', ps, ps)
	return G