def test_subs_lambda(): z = Variable('z', reals()) i = Variable('i', bint(5)) ix = random_tensor(OrderedDict([('i', bint(5))]), reals()) actual = Lambda(i, z)(z=ix) expected = Lambda(i(i='j'), z(z=ix)) check_funsor(actual, expected.inputs, expected.output) assert_close(actual, expected)
def test_slice_lambda(): z = Variable('z', reals()) i = Variable('i', bint(5)) j = Variable('j', bint(7)) zi = Lambda(i, z) zj = Lambda(j, z) zij = Lambda(j, zi) zj2 = zij[:, i] check_funsor(zj2, zj.inputs, zj.output)
def test_slice_lambda(): z = Variable('z', Real) i = Variable('i', Bint[5]) j = Variable('j', Bint[7]) zi = Lambda(i, z) zj = Lambda(j, z) zij = Lambda(j, zi) zj2 = zij[:, i] check_funsor(zj2, zj.inputs, zj.output)
def test_lambda_getitem(): data = randn((2, )) x = Tensor(data) y = Tensor(data, OrderedDict(i=bint(2))) i = Variable('i', bint(2)) assert x[i] is y assert Lambda(i, y) is x
def test_lambda(base_shape): z = Variable('z', reals(*base_shape)) i = Variable('i', bint(5)) j = Variable('j', bint(7)) zi = Lambda(i, z) assert zi.output.shape == (5, ) + base_shape assert zi[i] is z zj = Lambda(j, z) assert zj.output.shape == (7, ) + base_shape assert zj[j] is z zij = Lambda(j, zi) assert zij.output.shape == (7, 5) + base_shape assert zij[j] is zi assert zij[j, i] is z # assert zij[:, i] is zj # XXX this was disabled by alpha-renaming check_funsor(zij[:, i], zj.inputs, zj.output)
def eager_independent(delta, reals_var, bint_var): if delta.name == reals_var or delta.name.startswith(reals_var + "__BOUND"): i = Variable(bint_var, delta.inputs[bint_var]) point = Lambda(i, delta.point) if bint_var in delta.log_density.inputs: log_density = delta.log_density.reduce(ops.add, bint_var) else: log_density = delta.log_density * delta.inputs[bint_var].dtype return Delta(reals_var, point, log_density) return None # defer to default implementation
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 eager_independent_delta(delta, reals_var, bint_var, diag_var): for i, (name, (point, log_density)) in enumerate(delta.terms): if name == diag_var: bv = Variable(bint_var, delta.inputs[bint_var]) point = Lambda(bv, point) if bint_var in log_density.inputs: log_density = log_density.reduce(ops.add, bint_var) else: log_density = log_density * delta.inputs[bint_var].dtype new_terms = delta.terms[:i] + ((reals_var, (point, log_density)),) + delta.terms[i+1:] return Delta(new_terms) return None
def test_quote(interp): with interpretation(interp): x = Variable('x', bint(8)) check_quote(x) y = Variable('y', reals(8, 3, 3)) check_quote(y) check_quote(y[x]) z = Stack('i', (Number(0), Variable('z', reals()))) check_quote(z) check_quote(z(i=0)) check_quote(z(i=Slice('i', 0, 1, 1, 2))) check_quote(z.reduce(ops.add, 'i')) check_quote(Cat('i', (z, z, z))) check_quote(Lambda(Variable('i', bint(2)), z))
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