Cannot set posterior using RVPair
By specifying the second argument to RVPair, it should be possible to set posterior of a distribution used in a module.
For example (based on a test from borch / a conversation in Zulip):
module.weight = RVPair(
dist.Normal(0, 5),
dist.Normal(
Parameter(torch.tensor(2.0)),
borch.Transform(
torch.nn.functional.softplus, Parameter(torch.tensor(0.1))
),
),
)
Should set the value of the posterior for weight to 2.0. However, the value is unchanged. The following test reproduces the issue:
def test_setting_posterior():
model = borch.nn.Linear(4, 4)
ones = torch.ones(4, 4)
model.weight = borch.RVPair(
borch.distributions.Normal(ones, 5),
borch.distributions.Normal(
torch.nn.Parameter(ones),
borch.Transform(
torch.nn.functional.softplus,
torch.nn.Parameter(torch.tensor(0.1)),
),
),
)
assert model.prior.weight.loc.sum() == 16 # Works
assert model.posterior.weight.loc.sum() == 16 # Doesn't work
A quick fix is to manually set the posterior after using RVPair:
def test_setting_posterior():
model = borch.nn.Linear(4, 4)
ones = torch.ones(4, 4)
model.weight = borch.RVPair(
borch.distributions.Normal(ones, 5),
borch.distributions.Normal(
torch.nn.Parameter(ones),
borch.Transform(
torch.nn.functional.softplus,
torch.nn.Parameter(torch.tensor(0.1)),
),
),
)
model.posterior.weight = borch.distributions.Normal(
torch.nn.Parameter(ones),
borch.Transform(
torch.nn.functional.softplus,
torch.nn.Parameter(torch.tensor(0.1)),
),
)
assert model.prior.weight.loc.sum() == 16 # Works
assert model.posterior.weight.loc.sum() == 16 # Works
Code should be added in the relevant place(s) to get the expected behaviour without manually setting the posterior afterwards.