def eager_negate_variable(op, var): if var.dtype != "real": return None const = Number(0.) coeffs = ((var, Number(-1, "real")), ) return Affine(const, coeffs)
def model(size, position=0): if size == 1: name = str(position) return Uniform((Delta(name, Number(0, 2)), Delta(name, Number(1, 2)))) return Uniform( model(t, position) + model(size - t, t + position) for t in range(1, size))
def test_reduce_syntactic_sugar(): x = Stack("i", (Number(1), Number(2), Number(3))) expected = Number(1 + 2 + 3) assert x.reduce(ops.add) is expected assert x.reduce(ops.add, "i") is expected assert x.reduce(ops.add, {"i"}) is expected assert x.reduce(ops.add, frozenset(["i"])) is expected
def eager_multinomial(total_count, probs, value): # Multinomial.log_prob() supports inhomogeneous total_count only by # avoiding passing total_count to the constructor. inputs, (total_count, probs, value) = align_tensors(total_count, probs, value) shape = broadcast_shape(total_count.shape + (1,), probs.shape, value.shape) probs = Tensor(probs.expand(shape), inputs) value = Tensor(value.expand(shape), inputs) total_count = Number(total_count.max().item()) # Used by distributions validation code. return Multinomial.eager_log_prob(total_count=total_count, probs=probs, value=value)
def test_reduce_syntactic_sugar(): i = Variable("i", Bint[3]) x = Stack("i", (Number(1), Number(2), Number(3))) expected = Number(1 + 2 + 3) assert x.reduce(ops.add) is expected assert x.reduce(ops.add, "i") is expected assert x.reduce(ops.add, {"i"}) is expected assert x.reduce(ops.add, frozenset(["i"])) is expected assert x.reduce(ops.add, i) is expected assert x.reduce(ops.add, {i}) is expected assert x.reduce(ops.add, frozenset([i])) is expected
def test_unify_binary(): with interpretation(lazy): pattern = Variable('a', reals()) + Number(2.) * Variable('b', reals()) expr = Number(1.) + Number(2.) * (Number(3.) - Number(4.)) subs = unify(pattern, expr) print(subs, pattern(**{k.name: v for k, v in subs.items()})) assert subs is not False with interpretation(unify_interpreter): assert unify((pattern, ), (expr, )) is not False
def eager_binary(op, lhs, rhs): if lhs.dtype != "real" or rhs.dtype != "real": return None if op is ops.add: const = Number(0.) coeffs = ((lhs, Number(1.)), (rhs, Number(1.))) return Affine(const, coeffs) elif op is ops.sub: return lhs + -rhs return None
def test_slice(): t_slice = Slice("t", 10) s_slice = t_slice(t="s") assert isinstance(s_slice, Slice) assert s_slice.slice == t_slice.slice assert s_slice(s="t") is t_slice assert t_slice(t=0) is Number(0, 10) assert t_slice(t=1) is Number(1, 10) assert t_slice(t=2) is Number(2, 10) assert t_slice(t=t_slice) is t_slice
def test_binary(symbol, data1, data2): dtype = 'real' if symbol in BOOLEAN_OPS: dtype = 2 data1 = bool(data1) data2 = bool(data2) try: expected_data = binary_eval(symbol, data1, data2) except ZeroDivisionError: return x1 = Number(data1, dtype) x2 = Number(data2, dtype) actual = binary_eval(symbol, x1, x2) check_funsor(actual, {}, Domain((), dtype), expected_data)
def eager_reduce(self, op, reduced_vars): if op is ops.logaddexp: if reduced_vars - self.fresh and self.fresh - reduced_vars: result = self.eager_reduce(op, reduced_vars & self.fresh) if reduced_vars & self.fresh else self if result is not self: result = result.eager_reduce(op, reduced_vars - self.fresh) if reduced_vars - self.fresh else self return result if result is not self else None return None result_terms = [(name, (point, log_density)) for name, (point, log_density) in self.terms if name not in reduced_vars] result_terms, scale = [], Number(0) for name, (point, log_density) in self.terms: if name in reduced_vars: # XXX obscenely wasteful - need a lazy Zero term if point.inputs: scale += (point == point).all().log() if log_density.inputs: scale += log_density * 0. else: result_terms.append((name, (point, log_density))) result = Delta(tuple(result_terms)) + scale if result_terms else scale return result.reduce(op, reduced_vars - self.fresh) if op is ops.add: raise NotImplementedError("TODO Implement ops.add to simulate .to_event().") return None # defer to default implementation
def eager_subs(self, subs): terms = OrderedDict(self.terms) new_terms = terms.copy() log_density = Number(0) for name, value in subs: if isinstance(value, Variable): new_terms[value.name] = new_terms.pop(name) continue if not any(d.dtype == 'real' for side in (value, terms[name][0]) for d in side.inputs.values()): point, point_log_density = new_terms.pop(name) log_density += (value == point).all().log() + point_log_density continue # Try to invert the substitution. soln = solve(value, terms[name][0]) if soln is None: return None # lazily substitute new_name, new_point, point_log_density = soln old_point, old_point_density = new_terms.pop(name) new_terms[new_name] = (new_point, old_point_density + point_log_density) return Delta(tuple( new_terms.items())) + log_density if new_terms else log_density
def test_smoke(expr, expected_type): g1 = Gaussian(info_vec=numeric_array([[0.0, 0.1, 0.2], [2.0, 3.0, 4.0]]), precision=numeric_array([[[1.0, 0.1, 0.2], [0.1, 1.0, 0.3], [0.2, 0.3, 1.0]], [[1.0, 0.1, 0.2], [0.1, 1.0, 0.3], [0.2, 0.3, 1.0]]]), inputs=OrderedDict([('i', bint(2)), ('x', reals(3))])) assert isinstance(g1, Gaussian) g2 = Gaussian(info_vec=numeric_array([[0.0, 0.1], [2.0, 3.0]]), precision=numeric_array([[[1.0, 0.2], [0.2, 1.0]], [[1.0, 0.2], [0.2, 1.0]]]), inputs=OrderedDict([('i', bint(2)), ('y', reals(2))])) assert isinstance(g2, Gaussian) shift = Tensor(numeric_array([-1., 1.]), OrderedDict([('i', bint(2))])) assert isinstance(shift, Tensor) i0 = Number(1, 2) assert isinstance(i0, Number) x0 = Tensor(numeric_array([0.5, 0.6, 0.7])) assert isinstance(x0, Tensor) y0 = Tensor(numeric_array([[0.2, 0.3], [0.8, 0.9]]), inputs=OrderedDict([('i', bint(2))])) assert isinstance(y0, Tensor) result = eval(expr) assert isinstance(result, expected_type)
def test_smoke(expr, expected_type): dx = Delta('x', Tensor(torch.randn(2, 3), OrderedDict([('i', bint(2))]))) assert isinstance(dx, Delta) dy = Delta('y', Tensor(torch.randn(3, 4), OrderedDict([('j', bint(3))]))) assert isinstance(dy, Delta) t = Tensor(torch.randn(2, 3), OrderedDict([('i', bint(2)), ('j', bint(3))])) assert isinstance(t, Tensor) g = Gaussian(info_vec=torch.tensor([[0.0, 0.1, 0.2], [2.0, 3.0, 4.0]]), precision=torch.tensor([[[1.0, 0.1, 0.2], [0.1, 1.0, 0.3], [0.2, 0.3, 1.0]], [[1.0, 0.1, 0.2], [0.1, 1.0, 0.3], [0.2, 0.3, 1.0]]]), inputs=OrderedDict([('i', bint(2)), ('x', reals(3))])) assert isinstance(g, Gaussian) i0 = Number(1, 2) assert isinstance(i0, Number) x0 = Tensor(torch.tensor([0.5, 0.6, 0.7])) assert isinstance(x0, Tensor) result = eval(expr) assert isinstance(result, expected_type)
def eager_reduce(self, op, reduced_vars): if op is ops.logaddexp: if self.name in reduced_vars: return Number(0) # Deltas are normalized. # TODO Implement ops.add to simulate .to_event(). return None # defer to default implementation
def eager_binary(op, var, other): if var.dtype != "real" or other.dtype != "real": return None if op is ops.add: const = other coeffs = ((var, Number(1.)), ) return Affine(const, coeffs) elif op is ops.mul: const = Number(0.) coeffs = ((var, other), ) return Affine(const, coeffs) elif op is ops.sub: return var + -other elif op is ops.truediv: return var * (1. / other) return None
def __call__(cls, *args): if len(args) > 1: assert len(args) == 2 or len(args) == 3 assert isinstance(args[0], str) and isinstance(args[1], Funsor) args = args + (Number(0.), ) if len(args) == 2 else args args = (((args[0], (to_funsor(args[1]), to_funsor(args[2]))), ), ) assert isinstance(args[0], tuple) return super().__call__(args[0])
def eager_reduce(self, op, reduced_vars): if op is ops.logaddexp: # Keep mixture parameters lazy. mixture_vars = frozenset(k for k, d in self.gaussian.inputs.items() if d.dtype != 'real') mixture_vars = mixture_vars.union(*(x.point.inputs for x in self.deltas)) lazy_vars = reduced_vars & mixture_vars reduced_vars -= lazy_vars # Integrate out degenerate variables, i.e. drop selected delta. deltas = [] remaining_vars = set(reduced_vars) for d in self.deltas: if d.name in reduced_vars: remaining_vars.remove(d.name) else: deltas.append(d) deltas = tuple(deltas) reduced_vars = frozenset(remaining_vars) # Integrate out delayed discrete variables. discrete_vars = reduced_vars.intersection(self.discrete.inputs) discrete = self.discrete.reduce(op, discrete_vars) reduced_vars -= discrete_vars # Integrate out delayed gaussian variables. gaussian_vars = reduced_vars.intersection(self.gaussian.inputs) gaussian = self.gaussian.reduce(ops.logaddexp, gaussian_vars) reduced_vars -= gaussian_vars # Scale to account for remaining reduced_vars that were inputs to dropped deltas. eager_result = Joint(deltas, discrete) if gaussian is not Number(0): eager_result += gaussian reduced_vars |= lazy_vars.difference(eager_result.inputs) lazy_vars = lazy_vars.intersection(eager_result.inputs) if reduced_vars: eager_result += ops.log( reduce(ops.mul, [self.inputs[v].dtype for v in reduced_vars])) # Return a value only if progress has been made. if eager_result is self: return None # defer to default implementation else: return eager_result.reduce(ops.logaddexp, lazy_vars) if op is ops.add: terms = list(self.deltas) + [self.discrete, self.gaussian] for i, term in enumerate(terms): terms[i] = term.reduce(ops.add, reduced_vars.intersection(term.inputs)) return reduce(ops.add, terms) return None # defer to default implementation
def test_reduce_moment_matching_moments(): x = Variable('x', Reals[2]) gaussian = random_gaussian( OrderedDict([('i', Bint[2]), ('j', Bint[3]), ('x', Reals[2])])) with interpretation(moment_matching): approx = gaussian.reduce(ops.logaddexp, 'j') with interpretation(MonteCarlo(s=Bint[100000])): actual = Integrate(approx, Number(1.), 'x') expected = Integrate(gaussian, Number(1.), {'j', 'x'}) assert_close(actual, expected, atol=1e-3, rtol=1e-3) actual = Integrate(approx, x, 'x') expected = Integrate(gaussian, x, {'j', 'x'}) assert_close(actual, expected, atol=1e-2, rtol=1e-2) actual = Integrate(approx, x * x, 'x') expected = Integrate(gaussian, x * x, {'j', 'x'}) assert_close(actual, expected, atol=1e-2, rtol=1e-2)
def test_binary(symbol, data1, data2): dtype = 'real' if symbol in BOOLEAN_OPS: dtype = 2 data1 = bool(data1) data2 = bool(data2) try: expected_data = binary_eval(symbol, data1, data2) except ZeroDivisionError: return x1 = Number(data1, dtype) x2 = Number(data2, dtype) actual = binary_eval(symbol, x1, x2) check_funsor(actual, {}, Domain((), dtype), expected_data) with interpretation(normalize): actual_reflect = binary_eval(symbol, x1, x2) assert actual.output == actual_reflect.output
def test_reduce_moment_matching_moments(): x = Variable('x', reals(2)) gaussian = random_gaussian( OrderedDict([('i', bint(2)), ('j', bint(3)), ('x', reals(2))])) with interpretation(moment_matching): approx = gaussian.reduce(ops.logaddexp, 'j') with monte_carlo_interpretation(s=bint(100000)): actual = Integrate(approx, Number(1.), 'x') expected = Integrate(gaussian, Number(1.), {'j', 'x'}) assert_close(actual, expected, atol=1e-3, rtol=1e-3) actual = Integrate(approx, x, 'x') expected = Integrate(gaussian, x, {'j', 'x'}) assert_close(actual, expected, atol=1e-2, rtol=1e-2) actual = Integrate(approx, x * x, 'x') expected = Integrate(gaussian, x * x, {'j', 'x'}) assert_close(actual, expected, atol=1e-2, rtol=1e-2)
def sum_product(sum_op, prod_op, factors, eliminate=frozenset(), plates=frozenset()): """ Performs sum-product contraction of a collection of factors. :return: a single contracted Funsor. :rtype: :class:`~funsor.terms.Funsor` """ factors = partial_sum_product(sum_op, prod_op, factors, eliminate, plates) return reduce(prod_op, factors, Number(UNITS[prod_op]))
def test_advanced_indexing_lazy(output_shape): x = Tensor(randn((2, 3, 4) + output_shape), OrderedDict([ ('i', bint(2)), ('j', bint(3)), ('k', bint(4)), ])) u = Variable('u', bint(2)) v = Variable('v', bint(3)) with interpretation(lazy): i = Number(1, 2) - u j = Number(2, 3) - v k = u + v expected_data = empty((2, 3) + output_shape) i_data = x.materialize(i).data j_data = x.materialize(j).data k_data = x.materialize(k).data for u in range(2): for v in range(3): expected_data[u, v] = x.data[i_data[u], j_data[v], k_data[u, v]] expected = Tensor(expected_data, OrderedDict([ ('u', bint(2)), ('v', bint(3)), ])) assert_equiv(expected, x(i, j, k)) assert_equiv(expected, x(i=i, j=j, k=k)) assert_equiv(expected, x(i=i, j=j)(k=k)) assert_equiv(expected, x(j=j, k=k)(i=i)) assert_equiv(expected, x(k=k, i=i)(j=j)) assert_equiv(expected, x(i=i)(j=j, k=k)) assert_equiv(expected, x(j=j)(k=k, i=i)) assert_equiv(expected, x(k=k)(i=i, j=j)) assert_equiv(expected, x(i=i)(j=j)(k=k)) assert_equiv(expected, x(i=i)(k=k)(j=j)) assert_equiv(expected, x(j=j)(i=i)(k=k)) assert_equiv(expected, x(j=j)(k=k)(i=i)) assert_equiv(expected, x(k=k)(i=i)(j=j)) assert_equiv(expected, x(k=k)(j=j)(i=i))
def test_match_binary(): with interpretation(lazy): pattern = Variable('a', reals()) + Number(2.) * Variable('b', reals()) expr = Number(1.) + Number(2.) * (Number(3.) - Number(4.)) @match_vars(pattern) def expand_2_vars(a, b): return a + b + b @match(pattern) def expand_2_walk(x): return x.lhs + x.rhs.rhs + x.rhs.rhs eager_val = reinterpret(expr) lazy_val = expand_2_vars(expr) assert eager_val == reinterpret(lazy_val) lazy_val_2 = expand_2_walk(expr) assert eager_val == reinterpret(lazy_val_2)
def test_unary(symbol, data): dtype = 'real' if symbol == '~': data = bool(data) dtype = 2 expected_data = unary_eval(symbol, data) x = Number(data, dtype) actual = unary_eval(symbol, x) check_funsor(actual, {}, Domain((), dtype), expected_data)
def eager_add(op, joint, other): # Update with a delayed gaussian random variable. subs = tuple( (d.name, d.point) for d in joint.deltas if d.name in other.inputs) if subs: other = Subs(other, subs) if joint.gaussian is not Number(0): other = joint.gaussian + other if not isinstance(other, Gaussian): return Joint(joint.deltas, joint.discrete) + other return Joint(joint.deltas, joint.discrete, other)
def test_unary(symbol, data): dtype = 'real' if symbol == '~': data = bool(data) dtype = 2 if symbol == 'atanh': data = min(data, 0.99) expected_data = unary_eval(symbol, data) x = Number(data, dtype) actual = unary_eval(symbol, x) check_funsor(actual, {}, Array[dtype, ()], expected_data)
def eager_joint(deltas, discrete, gaussian): if not isinstance(gaussian, (Number, Tensor, Gaussian)): return Joint(deltas, discrete) + gaussian if any(not isinstance(d, Delta) for d in deltas): new_deltas = [] for d in deltas: if isinstance(d, Delta): new_deltas.append(d) elif isinstance(d, (Number, Tensor)): discrete += d else: raise ValueError("Invalid component for Joint: {}".format(d)) return Joint(tuple(new_deltas), discrete) + gaussian if isinstance(gaussian, (Number, Tensor)) and gaussian is not Number(0): discrete += gaussian return Joint(deltas, discrete, Number(0)) # Demote a Joint to a simpler elementary funsor. if not deltas: if gaussian is Number(0): return discrete elif discrete is Number(0): return gaussian elif len(deltas) == 1: if discrete is Number(0) and gaussian is Number(0): return deltas[0] return None # defer to default implementation
def test_cons_hash(): assert Variable('x', bint(3)) is Variable('x', bint(3)) assert Variable('x', reals()) is Variable('x', reals()) assert Variable('x', reals()) is not Variable('x', bint(3)) assert Number(0, 3) is Number(0, 3) assert Number(0.) is Number(0.) assert Number(0.) is not Number(0, 3)
def adjoint(expr, targets, start=Number(0.)): adjoint_values = defaultdict(lambda: Number(0.)) # 1 in logspace multiplicities = defaultdict(lambda: 0) tape_recorder = AdjointTape() with interpretation(tape_recorder): adjoint_values[reinterpret(expr)] = start while tape_recorder.tape: output, fn, inputs = tape_recorder.tape.pop() in_adjs = adjoint_ops(fn, adjoint_values[output], output, *inputs) for v, adjv in in_adjs.items(): multiplicities[v] += 1 adjoint_values[v] = adjoint_values[v] + adjv # product in logspace target_adjs = {} for v in targets: target_adjs[v] = adjoint_values[v] / multiplicities[v] if not isinstance(v, Variable): target_adjs[v] = target_adjs[v] + v return target_adjs
def eager_binary_affine_variable(op, affine, other): if op is ops.add: const = affine.const coeffs = affine.coeffs.copy() if other in affine.inputs: coeffs[other] += 1 else: coeffs[other] = Number(1.) return Affine(const, tuple(coeffs.items())) if op is ops.sub: return affine + -other return None