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
