def test_einsum(equation, backend):
inputs, outputs, sizes, operands, funsor_operands = make_einsum_example(
equation)
expected = opt_einsum.contract(equation, *operands, backend=backend)
with interpretation(reflect):
naive_ast = naive_einsum(equation, *funsor_operands, backend=backend)
optimized_ast = apply_optimizer(naive_ast)
print("Naive expression: {}".format(naive_ast))
print("Optimized expression: {}".format(optimized_ast))
actual_optimized = reinterpret(optimized_ast) # eager by default
actual = naive_einsum(equation, *funsor_operands, backend=backend)
assert isinstance(actual, funsor.Tensor) and len(outputs) == 1
if len(outputs[0]) > 0:
actual = actual.align(tuple(outputs[0]))
actual_optimized = actual_optimized.align(tuple(outputs[0]))
assert_close(actual, actual_optimized, atol=1e-4)
assert expected.shape == actual.data.shape
assert_close(expected, actual.data, rtol=1e-5, atol=1e-8)
for output in outputs:
for i, output_dim in enumerate(output):
assert output_dim in actual.inputs
assert actual.inputs[output_dim].dtype == sizes[output_dim]
def test_einsum_categorical(equation):
if get_backend() == "jax":
from funsor.jax.distributions import Categorical
else:
from funsor.torch.distributions import Categorical
inputs, outputs, sizes, operands, _ = make_einsum_example(equation)
operands = [ops.abs(operand) / ops.abs(operand).sum(-1)[..., None]
for operand in operands]
expected = opt_einsum.contract(equation, *operands,
backend=BACKEND_TO_EINSUM_BACKEND[get_backend()])
with interpretation(reflect):
funsor_operands = [
Categorical(probs=Tensor(
operand,
inputs=OrderedDict([(d, Bint[sizes[d]]) for d in inp[:-1]])
))(value=Variable(inp[-1], Bint[sizes[inp[-1]]])).exp()
for inp, operand in zip(inputs, operands)
]
naive_ast = naive_einsum(equation, *funsor_operands)
optimized_ast = apply_optimizer(naive_ast)
print("Naive expression: {}".format(naive_ast))
print("Optimized expression: {}".format(optimized_ast))
actual_optimized = reinterpret(optimized_ast) # eager by default
actual = naive_einsum(equation, *map(reinterpret, funsor_operands))
if len(outputs[0]) > 0:
actual = actual.align(tuple(outputs[0]))
actual_optimized = actual_optimized.align(tuple(outputs[0]))
assert_close(actual, actual_optimized, atol=1e-4)
assert expected.shape == actual.data.shape
assert_close(expected, actual.data)
for output in outputs:
for i, output_dim in enumerate(output):
assert output_dim in actual.inputs
assert actual.inputs[output_dim].dtype == sizes[output_dim]
def test_einsum_categorical(equation):
inputs, outputs, sizes, operands, _ = make_einsum_example(equation)
operands = [operand.abs() / operand.abs().sum(-1, keepdim=True)
for operand in operands]
expected = opt_einsum.contract(equation, *operands, backend='torch')
with interpretation(reflect):
funsor_operands = [
Categorical(probs=Tensor(
operand,
inputs=OrderedDict([(d, bint(sizes[d])) for d in inp[:-1]])
))(value=Variable(inp[-1], bint(sizes[inp[-1]]))).exp()
for inp, operand in zip(inputs, operands)
]
naive_ast = naive_einsum(equation, *funsor_operands)
optimized_ast = apply_optimizer(naive_ast)
print("Naive expression: {}".format(naive_ast))
print("Optimized expression: {}".format(optimized_ast))
actual_optimized = reinterpret(optimized_ast) # eager by default
actual = naive_einsum(equation, *map(reinterpret, funsor_operands))
if len(outputs[0]) > 0:
actual = actual.align(tuple(outputs[0]))
actual_optimized = actual_optimized.align(tuple(outputs[0]))
assert_close(actual, actual_optimized, atol=1e-4)
assert expected.shape == actual.data.shape
assert torch.allclose(expected, actual.data)
for output in outputs:
for i, output_dim in enumerate(output):
assert output_dim in actual.inputs
assert actual.inputs[output_dim].dtype == sizes[output_dim]