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