def eager_binary_tensor_tensor(op, lhs, rhs):
# Compute inputs and outputs.
dtype = find_domain(op, lhs.output, rhs.output).dtype
if lhs.inputs == rhs.inputs:
inputs = lhs.inputs
lhs_data, rhs_data = lhs.data, rhs.data
else:
inputs, (lhs_data, rhs_data) = align_tensors(lhs, rhs)
if len(lhs.shape) == 1:
lhs_data = ops.unsqueeze(lhs_data, -2)
if len(rhs.shape) == 1:
rhs_data = ops.unsqueeze(rhs_data, -1)
# Reshape to support broadcasting of output shape.
if inputs:
lhs_dim = max(2, len(lhs.shape))
rhs_dim = max(2, len(rhs.shape))
if lhs_dim < rhs_dim:
cut = len(lhs_data.shape) - lhs_dim
shape = lhs_data.shape
shape = shape[:cut] + (1, ) * (rhs_dim - lhs_dim) + shape[cut:]
lhs_data = lhs_data.reshape(shape)
elif rhs_dim < lhs_dim:
cut = len(rhs_data.shape) - rhs_dim
shape = rhs_data.shape
shape = shape[:cut] + (1, ) * (lhs_dim - rhs_dim) + shape[cut:]
rhs_data = rhs_data.reshape(shape)
data = op(lhs_data, rhs_data)
if len(lhs.shape) == 1:
data = data.squeeze(-2)
if len(rhs.shape) == 1:
data = data.squeeze(-1)
return Tensor(data, inputs, dtype)
def test_reduce_moment_matching_univariate():
int_inputs = [('i', bint(2))]
real_inputs = [('x', reals())]
inputs = OrderedDict(int_inputs + real_inputs)
int_inputs = OrderedDict(int_inputs)
real_inputs = OrderedDict(real_inputs)
p = 0.8
t = 1.234
s1, s2, s3 = 2.0, 3.0, 4.0
loc = numeric_array([[-s1], [s1]])
precision = numeric_array([[[s2**-2]], [[s3**-2]]])
info_vec = (precision @ ops.unsqueeze(loc, -1)).squeeze(-1)
discrete = Tensor(ops.log(numeric_array([1 - p, p])) + t, int_inputs)
gaussian = Gaussian(info_vec, precision, inputs)
gaussian -= gaussian.log_normalizer
joint = discrete + gaussian
with interpretation(moment_matching):
actual = joint.reduce(ops.logaddexp, 'i')
assert_close(actual.reduce(ops.logaddexp), joint.reduce(ops.logaddexp))
expected_loc = numeric_array([(2 * p - 1) * s1])
expected_variance = (4 * p * (1 - p) * s1**2 + (1 - p) * s2**2 + p * s3**2)
expected_precision = numeric_array([[1 / expected_variance]])
expected_info_vec = (
expected_precision @ ops.unsqueeze(expected_loc, -1)).squeeze(-1)
expected_gaussian = Gaussian(expected_info_vec, expected_precision,
real_inputs)
expected_gaussian -= expected_gaussian.log_normalizer
expected_discrete = Tensor(numeric_array(t))
expected = expected_discrete + expected_gaussian
assert_close(actual, expected, atol=1e-5, rtol=None)
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 == frozenset([integrand.name]):
loc = ops.cholesky_solve(ops.unsqueeze(log_measure.info_vec, -1), log_measure._precision_chol).squeeze(-1)
data = loc * ops.unsqueeze(ops.exp(log_measure.log_normalizer.data), -1)
data = data.reshape(loc.shape[:-1] + integrand.output.shape)
inputs = OrderedDict((k, d) for k, d in log_measure.inputs.items() if d.dtype != 'real')
result = Tensor(data, inputs)
return result.reduce(ops.add, reduced_vars - real_vars)
return None # defer to default implementation
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 = ops.exp(lhs.log_normalizer.data)
lhs_cov = ops.cholesky_inverse(lhs._precision_chol)
lhs_loc = ops.cholesky_solve(ops.unsqueeze(lhs.info_vec, -1), 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 unscaled_sample(self, sampled_vars, sample_inputs, rng_key=None):
sampled_vars = sampled_vars.intersection(self.inputs)
if not sampled_vars:
return self
if any(self.inputs[k].dtype != 'real' for k in sampled_vars):
raise ValueError(
'Sampling from non-normalized Gaussian mixtures is intentionally '
'not implemented. You probably want to normalize. To work around, '
'add a zero Tensor/Array with given inputs.')
# Partition inputs into sample_inputs + int_inputs + real_inputs.
sample_inputs = OrderedDict(
(k, d) for k, d in sample_inputs.items() if k not in self.inputs)
sample_shape = tuple(int(d.dtype) for d in sample_inputs.values())
int_inputs = OrderedDict(
(k, d) for k, d in self.inputs.items() if d.dtype != 'real')
real_inputs = OrderedDict(
(k, d) for k, d in self.inputs.items() if d.dtype == 'real')
inputs = sample_inputs.copy()
inputs.update(int_inputs)
if sampled_vars == frozenset(real_inputs):
shape = sample_shape + self.info_vec.shape
backend = get_backend()
if backend != "numpy":
from importlib import import_module
dist = import_module(funsor.distribution.
BACKEND_TO_DISTRIBUTIONS_BACKEND[backend])
sample_args = (shape, ) if rng_key is None else (rng_key,
shape)
white_noise = dist.Normal.dist_class(0, 1).sample(*sample_args)
else:
white_noise = np.random.randn(*shape)
white_noise = ops.unsqueeze(white_noise, -1)
white_vec = ops.triangular_solve(self.info_vec[..., None],
self._precision_chol)
sample = ops.triangular_solve(white_noise + white_vec,
self._precision_chol,
transpose=True)[..., 0]
offsets, _ = _compute_offsets(real_inputs)
results = []
for key, domain in real_inputs.items():
data = sample[...,
offsets[key]:offsets[key] + domain.num_elements]
data = data.reshape(shape[:-1] + domain.shape)
point = Tensor(data, inputs)
assert point.output == domain
results.append(Delta(key, point))
results.append(self.log_normalizer)
return reduce(ops.add, results)
raise NotImplementedError(
'TODO implement partial sampling of real variables')
def random_gaussian(inputs):
"""
Creates a random :class:`funsor.gaussian.Gaussian` with given inputs.
"""
assert isinstance(inputs, OrderedDict)
batch_shape = tuple(d.dtype for d in inputs.values() if d.dtype != 'real')
event_shape = (sum(d.num_elements for d in inputs.values()
if d.dtype == 'real'), )
prec_sqrt = randn(batch_shape + event_shape + event_shape)
precision = ops.matmul(prec_sqrt, ops.transpose(prec_sqrt, -1, -2))
precision = precision + 0.5 * ops.new_eye(precision, event_shape[:1])
loc = randn(batch_shape + event_shape)
info_vec = ops.matmul(precision, ops.unsqueeze(loc, -1)).squeeze(-1)
return Gaussian(info_vec, precision, inputs)
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 _parse_slices(index, value):
if not isinstance(index, tuple):
index = (index, )
if index[0] is Ellipsis:
index = index[1:]
start_stops = []
for pos, i in reversed(list(enumerate(index))):
if isinstance(i, slice):
start_stops.append((i.start, i.stop))
elif isinstance(i, int):
start_stops.append((i, i + 1))
value = ops.unsqueeze(value, pos - len(index))
else:
raise ValueError("invalid index: {}".format(i))
start_stops.reverse()
return start_stops, value
def _mv(mat, vec):
return ops.matmul(mat, ops.unsqueeze(vec, -1)).squeeze(-1)
def _vv(vec1, vec2):
"""
Computes the inner product ``< vec1 | vec 2 >``.
"""
return ops.matmul(ops.unsqueeze(vec1, -2),
ops.unsqueeze(vec2, -1)).squeeze(-1).squeeze(-1)
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
