def mean(self) -> NaturalMultivariateNormal:
mean, nu, a, B = self.to_standard()
expec_eta2 = -.5 * a * posdef_inverse(B)[0]
expec_eta1 = -2. * matvec(expec_eta2, mean)
return NaturalMultivariateNormal(expec_eta1,
expec_eta2,
validate_args=self._validate_args)
def _condition(self, y_dims, x_dims, x, squeeze):
"""Conditional distribution of Y|X"""
x_dims = torch.as_tensor(x_dims) # .unsqueeze(0)
y_dims = torch.as_tensor(y_dims) # .unsqueeze(0)
x = torch.cat(x, dim=-1)
# TODO: Correctly handle single dimensions (int)
m, S = self.loc, self.covariance_matrix
mx = m[..., x_dims] # (..., Dx)
my = m[..., y_dims] # (..., Dy)
Sxx = S[..., x_dims.unsqueeze(-1),
x_dims.unsqueeze(-2)] # (..., Dx, Dx)
Sxy = S[..., x_dims.unsqueeze(-1),
y_dims.unsqueeze(-2)] # (..., Dx, Dy)
Syy = S[..., y_dims.unsqueeze(-1),
y_dims.unsqueeze(-2)] # (..., Dy, Dy)
Syx = Sxy.transpose(-2, -1) # (..., Dy, Dx)
Syx_iSxx = Syx @ posdef_inverse(Sxx)[0] # (..., Dy, Dx)
# if x_dims.dim() > 0:
# Syx_iSxx = Syx @ posdef_inverse(Sxx)[0] # (..., Dy, Dx)
# else:
# Syx_iSxx = Syx / Sxx # (..., Dy, Dx)
Syy_x = Syy - Syx_iSxx @ Sxy # (..., Dy, Dy)
my_x = my + matvec(Syx_iSxx, x - mx) # (..., Dy)
if squeeze:
return td.Normal(my_x.squeeze(-1),
torch.sqrt(Syy_x).reshape(Syy_x.shape[:-2]))
return MultivariateNormal(my_x, covariance_matrix=Syy_x)
def from_standard(
mvn: td.MultivariateNormal) -> 'NaturalMultivariateNormal':
precision = mvn.precision_matrix
nat_param2 = -.5 * precision
nat_param1 = matvec(precision, mvn.mean)
return NaturalMultivariateNormal(nat_param1,
nat_param2,
validate_args=mvn._validate_args)
def expected_stats(self):
mean, nu, a, B = self.to_standard()
D = mean.shape[-1]
maha, B_tril = mahalanobis(B, mean)
logdet_B = 2. * triangular_logdet(B_tril)
expec_lambda2 = -.5 * a[..., None, None] * cholseky_inverse(B_tril)
expec_lambda1 = -2. * matvec(expec_lambda2, mean)
expec_log_norm = .5 * (a * maha + D / nu + mvdigamma(a, D) - logdet_B)
return expec_lambda1, expec_lambda2, expec_log_norm
def rsample(self, sample_shape=torch.Size()):
# X ~ Normal(0, I)
# Z ~ Chi2(df)
# Y = X / sqrt(Z / df) ~ MultivariateStudentT(df)
shape = self._extended_shape(sample_shape)
X = _standard_normal(shape, dtype=self.df.dtype, device=self.df.device)
Z = self._chi2.rsample(sample_shape)
Y = X * torch.rsqrt(Z / self.df).unsqueeze(-1)
return self.loc + matvec(self.scale_tril, Y)
def rsample(self, sample_shape=torch.Size()) -> NaturalMultivariateNormal:
# TODO: Test NW sampling
prior_mean, nu, a, B = self.to_standard()
prec = Wishart(a, B).rsample(sample_shape)
mean = td.MultivariateNormal(prior_mean, precision_matrix=nu *
prec).rsample(sample_shape)
eta2 = -.5 * prec
eta1 = matvec(prec, mean)
return NaturalMultivariateNormal(eta1,
eta2,
validate_args=self._validate_args)
def rsample(self, sample_shape=torch.Size()):
dtype, device = self.nat_param1.dtype, self.nat_param1.device
noise = torch.randn(*self._extended_shape(sample_shape),
dtype=dtype,
device=device)
return self.mean + matvec(self.scale_tril, noise)
def mean(self) -> torch.Tensor:
return matvec(self.covariance_matrix, self.nat_param1)
def log_prob(self, value: torch.Tensor) -> torch.Tensor:
if self._validate_args:
self._validate_sample(value)
prod1 = (value * self.nat_param1).sum(-1)
prod2 = (value * matvec(self.nat_param2, value)).sum(-1)
return prod1 + prod2 - self.log_normalizer