Affine Pattern Matching

affine_inputs(fn)[source]

Returns a [sound sub]set of real inputs of fn wrt which fn is known to be affine.

Parameters

fn (Funsor) – A funsor.

Returns

A set of input names wrt which fn is affine.

Return type

frozenset

extract_affine(fn)[source]

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 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 Memoize() interpretation.

Parameters

fn (Funsor) – A funsor that is affine wrt the (add,mul) semiring in some subset of its inputs.

Returns

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.

Return type

tuple

is_affine(fn)[source]

A sound but incomplete test to determine whether a funsor is affine with respect to all of its real inputs.

Parameters

fn (Funsor) – A funsor.

Return type

bool