Repository: aleximmer/Laplace (language: python)

Issue:
`FullSubnetLaplace.sample()` doesn't use `generator`
Hi

I was doing some reproducibility tests for `pred_type="nn"`, so testing `SomeLaplaceClass._nn_predictive_samples()` -> `SomeLaplaceClass.sample()`, using the current `main` (508843dbc37498e02ef421a573e7c546a90a9443).

I found that `FullSubnetLaplace.sample()` doesn't use the passed `generator` (`ruff check --select ARG laplace/subnetlaplace.py` also catches this). The reason is probably that the used `torch.distributions.MultivariateNormal` doesn't offer an API to pass in a `generator`.

I guess it would not be too difficult to rewrite this in the way other sample methods do things -- generate and re-use a set of standard Normal samples (because of #91) and go from there.

Thanks.


Discussion / hints:
Hi Steve, thanks for catching this.

I agree, the fix is super simple as you mentioned. Probably could also be as simple as reusing `super().sample()`.

I have to admit that we don't pay too much attention to the `SubnetLaplace` implementation anymore since the same can be done more intuitively by switching off the unwanted subset's grads. The benefit of the latter is that the Jacobian for the linearized predictive will only be computed over the selected subset. Meanwhile, the former computes the *full* Jacobian first and then slices it using the subnet indices.

In any case, happy to accept your PR, as always!

Fix patch:
diff --git a/laplace/subnetlaplace.py b/laplace/subnetlaplace.py
index b801037..01685b9 100644
--- a/laplace/subnetlaplace.py
+++ b/laplace/subnetlaplace.py
@@ -4,7 +4,6 @@
 
 import torch
 from torch import nn
-from torch.distributions import MultivariateNormal
 
 from laplace.baselaplace import DiagLaplace, FullLaplace, Likelihood, ParametricLaplace
 from laplace.curvature import EFInterface, GGNInterface
@@ -192,9 +191,10 @@ def _init_H(self) -> None:
     def sample(
         self, n_samples: int = 100, generator: torch.Generator | None = None
     ) -> torch.Tensor:
-        # sample only subnetwork parameters and set all other parameters to their MAP estimates
-        dist = MultivariateNormal(loc=self.mean_subnet, scale_tril=self.posterior_scale)
-        subnet_samples = dist.sample((n_samples,))
+        samples = torch.randn(
+            n_samples, self.n_params_subnet, device=self._device, generator=generator
+        )
+        subnet_samples = self.mean_subnet[None, ...] + samples @ self.posterior_scale
         return self.assemble_full_samples(subnet_samples)
 
 


Test patch:
diff --git a/tests/test_subnetlaplace.py b/tests/test_subnetlaplace.py
index 90ab55d..31d15d4 100644
--- a/tests/test_subnetlaplace.py
+++ b/tests/test_subnetlaplace.py
@@ -818,3 +818,56 @@ def test_subnet_marginal_likelihood(
 
     lap.fit(loader)
     lap.log_marginal_likelihood()
+
+
+@pytest.mark.parametrize(
+    "likelihood,hessian_structure", product(likelihoods, hessian_structures)
+)
+def test_sample(model, likelihood, hessian_structure, class_loader, reg_loader):
+    loader = class_loader if likelihood == "classification" else reg_loader
+
+    subnetmask = RandomSubnetMask(model=model, n_params_subnet=10)
+    subnetmask.select()
+
+    lap = Laplace(
+        model,
+        likelihood=likelihood,
+        subset_of_weights="subnetwork",
+        subnetwork_indices=subnetmask.indices,
+        hessian_structure=hessian_structure,
+    )
+    lap.fit(loader)
+
+    n_samples = 20
+    generator = torch.Generator()
+    generator.manual_seed(123)
+    generator_state = generator.get_state()
+
+    # Generate different samples.
+    for gen in [generator, None]:
+        samples_1 = lap.sample(n_samples=n_samples, generator=gen)
+        samples_2 = lap.sample(n_samples=n_samples, generator=gen)
+        assert not (samples_1 == samples_2).all()
+    assert samples_1.shape == (n_samples, model.n_params)
+
+    # Test that generator gets used.
+    generator.set_state(generator_state)
+    samples_1 = lap.sample(n_samples=n_samples, generator=generator)
+    generator.set_state(generator_state)
+    samples_2 = lap.sample(n_samples=n_samples, generator=generator)
+    assert (samples_1 == samples_2).all()
+
+    # Test that only model params indexed by subnetmask.indices vary across samples.
+    model_params = parameters_to_vector(model.parameters())
+    fixed_mask = torch.ones(model.n_params, dtype=bool)
+    fixed_mask[subnetmask.indices] = False
+    n_fixed = model.n_params - lap.n_params_subnet
+    assert n_fixed == fixed_mask.nonzero().shape[0]
+    assert samples_1[:, fixed_mask].shape == (n_samples, n_fixed)
+    assert (
+        samples_1[:, fixed_mask] == model_params[fixed_mask].repeat(n_samples, 1)
+    ).all()
+    assert (samples_1[:, ~fixed_mask] == samples_1[:, subnetmask.indices]).all()
+    assert not (
+        samples_1[:, ~fixed_mask] == model_params[~fixed_mask].repeat(n_samples, 1)
+    ).all()