def eager_integrate(log_measure, integrand, reduced_vars):
real_vars = frozenset(k for k in reduced_vars
if log_measure.inputs[k].dtype == 'real')
if real_vars:
lhs_reals = frozenset(k for k, d in log_measure.inputs.items()
if d.dtype == 'real')
rhs_reals = frozenset(k for k, d in integrand.inputs.items()
if d.dtype == 'real')
if lhs_reals == real_vars and rhs_reals <= real_vars:
inputs = OrderedDict((k, d) for t in (log_measure, integrand)
for k, d in t.inputs.items())
lhs_info_vec, lhs_precision = align_gaussian(inputs, log_measure)
rhs_info_vec, rhs_precision = align_gaussian(inputs, integrand)
lhs = Gaussian(lhs_info_vec, lhs_precision, inputs)
# Compute the expectation of a non-normalized quadratic form.
# See "The Matrix Cookbook" (November 15, 2012) ss. 8.2.2 eq. 380.
# http://www.math.uwaterloo.ca/~hwolkowi/matrixcookbook.pdf
norm = lhs.log_normalizer.data.exp()
lhs_cov = cholesky_inverse(lhs._precision_chol)
lhs_loc = lhs.info_vec.unsqueeze(-1).cholesky_solve(
lhs._precision_chol).squeeze(-1)
vmv_term = _vv(lhs_loc,
rhs_info_vec - 0.5 * _mv(rhs_precision, lhs_loc))
data = norm * (vmv_term - 0.5 * _trace_mm(rhs_precision, lhs_cov))
inputs = OrderedDict(
(k, d) for k, d in inputs.items() if k not in reduced_vars)
result = Tensor(data, inputs)
return result.reduce(ops.add, reduced_vars - real_vars)
raise NotImplementedError('TODO implement partial integration')
return None # defer to default implementation
def eager_cat_homogeneous(name, part_name, *parts):
assert parts
output = parts[0].output
inputs = OrderedDict([(part_name, None)])
for part in parts:
assert part.output == output
assert part_name in part.inputs
inputs.update(part.inputs)
int_inputs = OrderedDict(
(k, v) for k, v in inputs.items() if v.dtype != "real")
real_inputs = OrderedDict(
(k, v) for k, v in inputs.items() if v.dtype == "real")
inputs = int_inputs.copy()
inputs.update(real_inputs)
discretes = []
info_vecs = []
precisions = []
for part in parts:
inputs[part_name] = part.inputs[part_name]
int_inputs[part_name] = inputs[part_name]
shape = tuple(d.size for d in int_inputs.values())
if isinstance(part, Gaussian):
discrete = None
gaussian = part
elif issubclass(type(part), GaussianMixture
): # TODO figure out why isinstance isn't working
discrete, gaussian = part.terms[0], part.terms[1]
discrete = ops.expand(align_tensor(int_inputs, discrete), shape)
else:
raise NotImplementedError("TODO")
discretes.append(discrete)
info_vec, precision = align_gaussian(inputs, gaussian)
info_vecs.append(ops.expand(info_vec, shape + (-1, )))
precisions.append(ops.expand(precision, shape + (-1, -1)))
if part_name != name:
del inputs[part_name]
del int_inputs[part_name]
dim = 0
info_vec = ops.cat(dim, *info_vecs)
precision = ops.cat(dim, *precisions)
inputs[name] = Bint[info_vec.shape[dim]]
int_inputs[name] = inputs[name]
result = Gaussian(info_vec, precision, inputs)
if any(d is not None for d in discretes):
for i, d in enumerate(discretes):
if d is None:
discretes[i] = ops.new_zeros(info_vecs[i],
info_vecs[i].shape[:-1])
discrete = ops.cat(dim, *discretes)
result = result + Tensor(discrete, int_inputs)
return result
def adjoint_subs_gaussianmixture_gaussianmixture(adj_redop, adj_binop, out_adj, arg, subs):
if any(v.dtype == 'real' and not isinstance(v, Variable) for k, v in subs):
raise NotImplementedError("TODO implement adjoint for substitution into Gaussian real variable")
# invert renaming
renames = tuple((v.name, k) for k, v in subs if isinstance(v, Variable))
out_adj = Subs(out_adj, renames)
# inverting advanced indexing
slices = tuple((k, v) for k, v in subs if not isinstance(v, Variable))
assert len(slices + renames) == len(subs)
in_adj_discrete = adjoint_ops(Subs, adj_redop, adj_binop, out_adj.terms[0], arg.terms[0], subs)[arg.terms[0]]
arg_int_inputs = OrderedDict((k, v) for k, v in arg.inputs.items() if v.dtype != 'real')
out_adj_int_inputs = OrderedDict((k, v) for k, v in out_adj.inputs.items() if v.dtype != 'real')
arg_real_inputs = OrderedDict((k, v) for k, v in arg.inputs.items() if v.dtype == 'real')
align_inputs = OrderedDict((k, v) for k, v in out_adj.terms[1].inputs.items() if v.dtype != 'real')
align_inputs.update(arg_real_inputs)
out_adj_info_vec, out_adj_precision = align_gaussian(align_inputs, out_adj.terms[1])
in_adj_info_vec = list(adjoint_ops(Subs, adj_redop, adj_binop, # ops.add, ops.mul,
Tensor(out_adj_info_vec, out_adj_int_inputs),
Tensor(arg.terms[1].info_vec, arg_int_inputs),
slices).values())[0]
in_adj_precision = list(adjoint_ops(Subs, adj_redop, adj_binop, # ops.add, ops.mul,
Tensor(out_adj_precision, out_adj_int_inputs),
Tensor(arg.terms[1].precision, arg_int_inputs),
slices).values())[0]
assert isinstance(in_adj_info_vec, Tensor)
assert isinstance(in_adj_precision, Tensor)
in_adj_gaussian = Gaussian(in_adj_info_vec.data, in_adj_precision.data, arg.inputs.copy())
in_adj = in_adj_gaussian + in_adj_discrete
return {arg: in_adj}