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})