def gaussian_to_data(funsor_dist, name_to_dim=None, normalized=False):
if normalized:
return to_data(funsor_dist.log_normalizer + funsor_dist, name_to_dim=name_to_dim)
loc = funsor_dist.info_vec.unsqueeze(-1).cholesky_solve(cholesky(funsor_dist.precision)).squeeze(-1)
int_inputs = OrderedDict((k, d) for k, d in funsor_dist.inputs.items() if d.dtype != "real")
loc = to_data(Tensor(loc, int_inputs), name_to_dim)
precision = to_data(Tensor(funsor_dist.precision, int_inputs), name_to_dim)
return dist.MultivariateNormal(loc, precision_matrix=precision)
def gaussian_to_data(funsor_dist, name_to_dim=None, normalized=False):
if normalized:
return to_data(funsor_dist.log_normalizer + funsor_dist,
name_to_dim=name_to_dim)
loc = ops.cholesky_solve(ops.unsqueeze(funsor_dist.info_vec, -1),
ops.cholesky(funsor_dist.precision)).squeeze(-1)
int_inputs = OrderedDict(
(k, d) for k, d in funsor_dist.inputs.items() if d.dtype != "real")
loc = to_data(Tensor(loc, int_inputs), name_to_dim)
precision = to_data(Tensor(funsor_dist.precision, int_inputs), name_to_dim)
backend_dist = import_module(
BACKEND_TO_DISTRIBUTIONS_BACKEND[get_backend()])
return backend_dist.MultivariateNormal.dist_class(
loc, precision_matrix=precision)
def eager_reduce(self, op, reduced_vars):
if op is ops.logaddexp:
# Marginalize out real variables, but keep mixtures lazy.
assert all(v in self.inputs for v in reduced_vars)
real_vars = frozenset(k for k, d in self.inputs.items()
if d.dtype == "real")
reduced_reals = reduced_vars & real_vars
reduced_ints = reduced_vars - real_vars
if not reduced_reals:
return None # defer to default implementation
inputs = OrderedDict((k, d) for k, d in self.inputs.items()
if k not in reduced_reals)
if reduced_reals == real_vars:
result = self.log_normalizer
else:
int_inputs = OrderedDict(
(k, v) for k, v in inputs.items() if v.dtype != 'real')
offsets, _ = _compute_offsets(self.inputs)
a = []
b = []
for key, domain in self.inputs.items():
if domain.dtype == 'real':
block = ops.new_arange(
self.info_vec, offsets[key],
offsets[key] + domain.num_elements, 1)
(b if key in reduced_vars else a).append(block)
a = ops.cat(-1, *a)
b = ops.cat(-1, *b)
prec_aa = self.precision[..., a[..., None], a]
prec_ba = self.precision[..., b[..., None], a]
prec_bb = self.precision[..., b[..., None], b]
prec_b = ops.cholesky(prec_bb)
prec_a = ops.triangular_solve(prec_ba, prec_b)
prec_at = ops.transpose(prec_a, -1, -2)
precision = prec_aa - ops.matmul(prec_at, prec_a)
info_a = self.info_vec[..., a]
info_b = self.info_vec[..., b]
b_tmp = ops.triangular_solve(info_b[..., None], prec_b)
info_vec = info_a - ops.matmul(prec_at, b_tmp)[..., 0]
log_prob = Tensor(
0.5 * len(b) * math.log(2 * math.pi) -
_log_det_tri(prec_b) + 0.5 * (b_tmp[..., 0]**2).sum(-1),
int_inputs)
result = log_prob + Gaussian(info_vec, precision, inputs)
return result.reduce(ops.logaddexp, reduced_ints)
elif op is ops.add:
for v in reduced_vars:
if self.inputs[v].dtype == 'real':
raise ValueError(
"Cannot sum along a real dimension: {}".format(
repr(v)))
# Fuse Gaussians along a plate. Compare to eager_add_gaussian_gaussian().
old_ints = OrderedDict(
(k, v) for k, v in self.inputs.items() if v.dtype != 'real')
new_ints = OrderedDict(
(k, v) for k, v in old_ints.items() if k not in reduced_vars)
inputs = OrderedDict((k, v) for k, v in self.inputs.items()
if k not in reduced_vars)
info_vec = Tensor(self.info_vec,
old_ints).reduce(ops.add, reduced_vars)
precision = Tensor(self.precision,
old_ints).reduce(ops.add, reduced_vars)
assert info_vec.inputs == new_ints
assert precision.inputs == new_ints
return Gaussian(info_vec.data, precision.data, inputs)
return None # defer to default implementation
def _precision_chol(self):
return ops.cholesky(self.precision)
def moment_matching_contract_joint(red_op, bin_op, reduced_vars, discrete,
gaussian):
approx_vars = frozenset(
k for k in reduced_vars
if k in gaussian.inputs and gaussian.inputs[k].dtype != 'real')
exact_vars = reduced_vars - approx_vars
if exact_vars and approx_vars:
return Contraction(red_op, bin_op, exact_vars, discrete,
gaussian).reduce(red_op, approx_vars)
if approx_vars and not exact_vars:
discrete += gaussian.log_normalizer
new_discrete = discrete.reduce(
ops.logaddexp, approx_vars.intersection(discrete.inputs))
new_discrete = discrete.reduce(
ops.logaddexp, approx_vars.intersection(discrete.inputs))
num_elements = reduce(ops.mul, [
gaussian.inputs[k].num_elements
for k in approx_vars.difference(discrete.inputs)
], 1)
if num_elements != 1:
new_discrete -= math.log(num_elements)
int_inputs = OrderedDict(
(k, d) for k, d in gaussian.inputs.items() if d.dtype != 'real')
probs = (discrete - new_discrete.clamp_finite()).exp()
old_loc = Tensor(
ops.cholesky_solve(ops.unsqueeze(gaussian.info_vec, -1),
gaussian._precision_chol).squeeze(-1),
int_inputs)
new_loc = (probs * old_loc).reduce(ops.add, approx_vars)
old_cov = Tensor(ops.cholesky_inverse(gaussian._precision_chol),
int_inputs)
diff = old_loc - new_loc
outers = Tensor(
ops.unsqueeze(diff.data, -1) * ops.unsqueeze(diff.data, -2),
diff.inputs)
new_cov = ((probs * old_cov).reduce(ops.add, approx_vars) +
(probs * outers).reduce(ops.add, approx_vars))
# Numerically stabilize by adding bogus precision to empty components.
total = probs.reduce(ops.add, approx_vars)
mask = ops.unsqueeze(ops.unsqueeze((total.data == 0), -1), -1)
new_cov.data = new_cov.data + mask * ops.new_eye(
new_cov.data, new_cov.data.shape[-1:])
new_precision = Tensor(
ops.cholesky_inverse(ops.cholesky(new_cov.data)), new_cov.inputs)
new_info_vec = (
new_precision.data @ ops.unsqueeze(new_loc.data, -1)).squeeze(-1)
new_inputs = new_loc.inputs.copy()
new_inputs.update(
(k, d) for k, d in gaussian.inputs.items() if d.dtype == 'real')
new_gaussian = Gaussian(new_info_vec, new_precision.data, new_inputs)
new_discrete -= new_gaussian.log_normalizer
return new_discrete + new_gaussian
return None