def eager_affine_normal(matrix, loc, scale, value_x, value_y):
assert len(matrix.output.shape) == 2
assert value_x.output == reals(matrix.output.shape[0])
assert value_y.output == reals(matrix.output.shape[1])
loc += value_x @ matrix
int_inputs, (loc, scale) = align_tensors(loc, scale, expand=True)
i_name = gensym("i")
y_name = gensym("y")
y_i_name = gensym("y_i")
int_inputs[i_name] = bint(value_y.output.shape[0])
loc = Tensor(loc, int_inputs)
scale = Tensor(scale, int_inputs)
y_dist = Independent(Normal(loc, scale, y_i_name), y_name, i_name, y_i_name)
return y_dist(**{y_name: value_y})
def alpha_convert(expr):
alpha_subs = {
name: interpreter.gensym(name + "__BOUND")
for name in expr.bound if "__BOUND" not in name
}
if not alpha_subs:
return expr
new_values = []
for v in expr._ast_values:
v = substitute(v, alpha_subs)
if isinstance(v, str) and v not in expr.fresh:
v = alpha_subs.get(v, v)
elif isinstance(v, frozenset):
swapped = v & frozenset(alpha_subs.keys())
v |= frozenset(alpha_subs[k] for k in swapped)
v -= swapped
elif isinstance(v, tuple) and isinstance(v[0], tuple) and len(v[0]) == 2 and \
isinstance(v[0][0], str) and isinstance(v[0][1], Funsor):
v = tuple(
(alpha_subs[k] if k in alpha_subs else k, vv) for k, vv in v)
elif isinstance(
v, OrderedDict): # XXX is this case ever actually triggered?
v = OrderedDict([(alpha_subs[k] if k in alpha_subs else k, vv)
for k, vv in v.items()])
new_values.append(v)
return reflect(type(expr), *new_values)
def test_distribute_reduce(lhs_vars, rhs_vars):
lhs_vars, rhs_vars = frozenset(lhs_vars), frozenset(rhs_vars)
lhs = random_tensor(OrderedDict([('i', bint(3)), ('j', bint(2))]), reals())
rhs = random_tensor(OrderedDict([('i', bint(3)), ('j', bint(2))]), reals())
with interpretation(reflect):
actual_lhs = lhs.reduce(ops.add, lhs_vars) if lhs_vars else lhs
actual_rhs = rhs.reduce(ops.add, rhs_vars) if rhs_vars else rhs
actual = actual_lhs * actual_rhs
lhs_subs = {v: gensym(v) for v in lhs_vars}
rhs_subs = {v: gensym(v) for v in rhs_vars}
expected = (lhs(**lhs_subs) * rhs(**rhs_subs)).reduce(
ops.add, frozenset(lhs_subs.values()) | frozenset(rhs_subs.values()))
assert_close(actual, expected)
def extract_affine(fn):
"""
Extracts an affine representation of a funsor, satisfying::
x = ...
const, coeffs = extract_affine(x)
y = sum(Einsum(eqn, (coeff, Variable(var, coeff.output)))
for var, (coeff, eqn) in coeffs.items())
assert_close(y, x)
assert frozenset(coeffs) == affine_inputs(x)
The ``coeffs`` will have one key per input wrt which ``fn`` is known to be
affine (via :func:`affine_inputs` ), and ``const`` and ``coeffs.values``
will all be constant wrt these inputs.
The affine approximation is computed by ev evaluating ``fn`` at
zero and each basis vector. To improve performance, users may want to run
under the :func:`~funsor.memoize.memoize` interpretation.
:param Funsor fn: A funsor that is affine wrt the (add,mul) semiring in
some subset of its inputs.
:return: A pair ``(const, coeffs)`` where const is a funsor with no real
inputs and ``coeffs`` is an OrderedDict mapping input name to a
``(coefficient, eqn)`` pair in einsum form.
:rtype: tuple
"""
# NB: this depends on the global default backend.
prototype = get_default_prototype()
# Determine constant part by evaluating fn at zero.
inputs = affine_inputs(fn)
inputs = OrderedDict((k, v) for k, v in fn.inputs.items() if k in inputs)
zeros = {
k: Tensor(ops.new_zeros(prototype, v.shape))
for k, v in inputs.items()
}
const = fn(**zeros)
# Determine linear coefficients by evaluating fn on basis vectors.
name = gensym('probe')
coeffs = OrderedDict()
for k, v in inputs.items():
dim = v.num_elements
var = Variable(name, bint(dim))
subs = zeros.copy()
subs[k] = Tensor(
ops.new_eye(prototype, (dim, )).reshape((dim, ) + v.shape))[var]
coeff = Lambda(var, fn(**subs) - const).reshape(v.shape + const.shape)
inputs1 = ''.join(map(opt_einsum.get_symbol, range(len(coeff.shape))))
inputs2 = inputs1[:len(v.shape)]
output = inputs1[len(v.shape):]
eqn = f'{inputs1},{inputs2}->{output}'
coeffs[k] = coeff, eqn
return const, coeffs
def _alpha_mangle(expr):
"""
Rename bound variables in expr to avoid conflict with any free variables.
FIXME this does not avoid conflict with other bound variables.
"""
alpha_subs = {name: interpreter.gensym(name + "__BOUND")
for name in expr.bound if "__BOUND" not in name}
if not alpha_subs:
return expr
ast_values = expr._alpha_convert(alpha_subs)
return reflect(type(expr), *ast_values)
def eager_normal(loc, scale, value):
assert loc.output == Real
assert scale.output == Real
assert value.output == Real
if not is_affine(loc) or not is_affine(value):
return None # lazy
info_vec = ops.new_zeros(scale.data, scale.data.shape + (1, ))
precision = ops.pow(scale.data, -2).reshape(scale.data.shape + (1, 1))
log_prob = -0.5 * math.log(2 * math.pi) - ops.log(scale).sum()
inputs = scale.inputs.copy()
var = gensym('value')
inputs[var] = Real
gaussian = log_prob + Gaussian(info_vec, precision, inputs)
return gaussian(**{var: value - loc})
def eager_mvn(loc, scale_tril, value):
assert len(loc.shape) == 1
assert len(scale_tril.shape) == 2
assert value.output == loc.output
if not is_affine(loc) or not is_affine(value):
return None # lazy
info_vec = scale_tril.data.new_zeros(scale_tril.data.shape[:-1])
precision = ops.cholesky_inverse(scale_tril.data)
scale_diag = Tensor(scale_tril.data.diagonal(dim1=-1, dim2=-2), scale_tril.inputs)
log_prob = -0.5 * scale_diag.shape[0] * math.log(2 * math.pi) - scale_diag.log().sum()
inputs = scale_tril.inputs.copy()
var = gensym('value')
inputs[var] = reals(scale_diag.shape[0])
gaussian = log_prob + Gaussian(info_vec, precision, inputs)
return gaussian(**{var: value - loc})
def extract_affine(fn):
"""
Extracts an affine representation of a funsor, which is exact for affine
funsors and approximate otherwise. For affine funsors this satisfies::
x = ...
const, coeffs = extract_affine(x)
y = sum(Einsum(eqn, (coeff, Variable(var, coeff.output)))
for var, (coeff, eqn) in coeffs.items())
assert_close(y, x)
The affine approximation is computed by ev evaluating ``fn`` at
zero and each basis vector. To improve performance, users may want to run
under the :func:`~funsor.memoize.memoize` interpretation.
:param Funsor fn: A funsor assumed to be affine wrt the (add,mul) semiring.
The affine assumption is not checked.
:return: A pair ``(const, coeffs)`` where const is a funsor with no real
inputs and ``coeffs`` is an OrderedDict mapping input name to a
``(coefficient, eqn)`` pair in einsum form.
:rtype: tuple
"""
# Determine constant part by evaluating fn at zero.
real_inputs = OrderedDict(
(k, v) for k, v in fn.inputs.items() if v.dtype == 'real')
zeros = {k: Tensor(torch.zeros(v.shape)) for k, v in real_inputs.items()}
const = fn(**zeros)
# Determine linear coefficients by evaluating fn on basis vectors.
name = gensym('probe')
coeffs = OrderedDict()
for k, v in real_inputs.items():
dim = v.num_elements
var = Variable(name, bint(dim))
subs = zeros.copy()
subs[k] = Tensor(torch.eye(dim).reshape((dim, ) + v.shape))[var]
coeff = Lambda(var, fn(**subs) - const).reshape(v.shape + const.shape)
inputs1 = ''.join(map(opt_einsum.get_symbol, range(len(coeff.shape))))
inputs2 = inputs1[:len(v.shape)]
output = inputs1[len(v.shape):]
eqn = f'{inputs1},{inputs2}->{output}'
coeffs[k] = coeff, eqn
return const, coeffs
def eager_affine_normal(matrix, loc, scale, value_x, value_y):
assert len(matrix.output.shape) == 2
assert value_x.output == reals(matrix.output.shape[0])
assert value_y.output == reals(matrix.output.shape[1])
tensors = (matrix, loc, scale, value_y)
int_inputs, tensors = align_tensors(*tensors)
matrix, loc, scale, value_y = tensors
assert value_y.size(-1) == loc.size(-1)
prec_sqrt = matrix / scale.unsqueeze(-2)
precision = prec_sqrt.matmul(prec_sqrt.transpose(-1, -2))
delta = (value_y - loc) / scale
info_vec = prec_sqrt.matmul(delta.unsqueeze(-1)).squeeze(-1)
log_normalizer = (-0.5 * loc.size(-1) * math.log(2 * math.pi)
- 0.5 * delta.pow(2).sum(-1) - scale.log().sum(-1))
precision = precision.expand(info_vec.shape + (-1,))
log_normalizer = log_normalizer.expand(info_vec.shape[:-1])
inputs = int_inputs.copy()
x_name = gensym("x")
inputs[x_name] = value_x.output
x_dist = Tensor(log_normalizer, int_inputs) + Gaussian(info_vec, precision, inputs)
return x_dist(**{x_name: value_x})
