def test_reduce_subset(dims, reduced_vars, op): reduced_vars = frozenset(reduced_vars) sizes = {'a': 3, 'b': 4, 'c': 5} shape = tuple(sizes[d] for d in dims) inputs = OrderedDict((d, bint(sizes[d])) for d in dims) data = rand(shape) + 0.5 dtype = 'real' if op in [ops.and_, ops.or_]: data = ops.astype(data, 'uint8') dtype = 2 x = Tensor(data, inputs, dtype) actual = x.reduce(op, reduced_vars) expected_inputs = OrderedDict( (d, bint(sizes[d])) for d in dims if d not in reduced_vars) reduced_vars &= frozenset(dims) if not reduced_vars: assert actual is x else: if reduced_vars == frozenset(dims): data = REDUCE_OP_TO_NUMERIC[op](data, None) else: for pos in reversed(sorted(map(dims.index, reduced_vars))): data = REDUCE_OP_TO_NUMERIC[op](data, pos) check_funsor(actual, expected_inputs, Domain((), dtype)) assert_close(actual, Tensor(data, expected_inputs, dtype), atol=1e-5, rtol=1e-5)
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 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 test_reduce_all(dims, op): sizes = {'a': 3, 'b': 4, 'c': 5} shape = tuple(sizes[d] for d in dims) inputs = OrderedDict((d, bint(sizes[d])) for d in dims) data = rand(shape) + 0.5 if op in [ops.and_, ops.or_]: data = ops.astype(data, 'uint8') expected_data = REDUCE_OP_TO_NUMERIC[op](data, None) x = Tensor(data, inputs) actual = x.reduce(op) check_funsor(actual, {}, reals(), expected_data)
def _get_stat_diff(funsor_dist_class, sample_inputs, inputs, num_samples, statistic, with_lazy, params): params = [Tensor(p, inputs) for p in params] if isinstance(with_lazy, bool): with interpretation(lazy if with_lazy else eager): funsor_dist = funsor_dist_class(*params) else: funsor_dist = funsor_dist_class(*params) rng_key = None if get_backend() == "torch" else np.array([0, 0], dtype=np.uint32) sample_value = funsor_dist.sample(frozenset(['value']), sample_inputs, rng_key=rng_key) expected_inputs = OrderedDict( tuple(sample_inputs.items()) + tuple(inputs.items()) + (('value', funsor_dist.inputs['value']), )) check_funsor(sample_value, expected_inputs, reals()) if sample_inputs: actual_mean = Integrate(sample_value, Variable('value', funsor_dist.inputs['value']), frozenset(['value' ])).reduce(ops.add, frozenset(sample_inputs)) inputs, tensors = align_tensors( *list(funsor_dist.params.values())[:-1]) raw_dist = funsor_dist.dist_class( **dict(zip(funsor_dist._ast_fields[:-1], tensors))) expected_mean = Tensor(raw_dist.mean, inputs) if statistic == "mean": actual_stat, expected_stat = actual_mean, expected_mean elif statistic == "variance": actual_stat = Integrate( sample_value, (Variable('value', funsor_dist.inputs['value']) - actual_mean)**2, frozenset(['value'])).reduce(ops.add, frozenset(sample_inputs)) expected_stat = Tensor(raw_dist.variance, inputs) elif statistic == "entropy": actual_stat = -Integrate(sample_value, funsor_dist, frozenset(['value'])).reduce( ops.add, frozenset(sample_inputs)) expected_stat = Tensor(raw_dist.entropy(), inputs) else: raise ValueError("invalid test statistic") diff = actual_stat.reduce(ops.add).data - expected_stat.reduce( ops.add).data return diff.sum(), diff
def _check_sample(funsor_dist, sample_inputs, inputs, atol=1e-2, rtol=None, num_samples=100000, statistic="mean", skip_grad=False): """utility that compares a Monte Carlo estimate of a distribution mean with the true mean""" samples_per_dim = int(num_samples**(1. / max(1, len(sample_inputs)))) sample_inputs = OrderedDict( (k, bint(samples_per_dim)) for k in sample_inputs) for tensor in list(funsor_dist.params.values())[:-1]: tensor.data.requires_grad_() sample_value = funsor_dist.sample(frozenset(['value']), sample_inputs) expected_inputs = OrderedDict( tuple(sample_inputs.items()) + tuple(inputs.items()) + (('value', funsor_dist.inputs['value']), )) check_funsor(sample_value, expected_inputs, reals()) if sample_inputs: actual_mean = Integrate(sample_value, Variable('value', funsor_dist.inputs['value']), frozenset(['value' ])).reduce(ops.add, frozenset(sample_inputs)) inputs, tensors = align_tensors( *list(funsor_dist.params.values())[:-1]) raw_dist = funsor_dist.dist_class( **dict(zip(funsor_dist._ast_fields[:-1], tensors))) expected_mean = Tensor(raw_dist.mean, inputs) check_funsor(actual_mean, expected_mean.inputs, expected_mean.output) assert_close(actual_mean, expected_mean, atol=atol, rtol=rtol) if sample_inputs and not skip_grad: if statistic == "mean": actual_stat, expected_stat = actual_mean, expected_mean elif statistic == "variance": actual_stat = Integrate( sample_value, (Variable('value', funsor_dist.inputs['value']) - actual_mean)**2, frozenset(['value'])).reduce(ops.add, frozenset(sample_inputs)) expected_stat = Tensor(raw_dist.variance, inputs) elif statistic == "entropy": actual_stat = -Integrate(sample_value, funsor_dist, frozenset(['value'])).reduce( ops.add, frozenset(sample_inputs)) expected_stat = Tensor(raw_dist.entropy(), inputs) else: raise ValueError("invalid test statistic") grad_targets = [v.data for v in list(funsor_dist.params.values())[:-1]] actual_grads = torch.autograd.grad(actual_stat.reduce( ops.add).sum().data, grad_targets, allow_unused=True) expected_grads = torch.autograd.grad(expected_stat.reduce( ops.add).sum().data, grad_targets, allow_unused=True) assert_close(actual_stat, expected_stat, atol=atol, rtol=rtol) for actual_grad, expected_grad in zip(actual_grads, expected_grads): if expected_grad is not None: assert_close(actual_grad, expected_grad, atol=atol, rtol=rtol) else: assert_close(actual_grad, torch.zeros_like(actual_grad), atol=atol, rtol=rtol)