Funsor Factory

class Fresh(fn)[source]

Bases: object

Type hint for make_funsor() decorated functions. This provides hints for fresh variables (names) and the return type.

Examples:

Fresh[Real]  # a constant known domain
Fresh[lambda x: Array[x.dtype, x.shape[1:]]  # args are Domains
Fresh[lambda x, y: Bint[x.size + y.size]]
Parameters

fn (callable) – A lambda taking named arguments (in any order) which will be filled in with the domain of the similarly named funsor argument to the decorated function. This lambda should compute a desired resulting domain given domains of arguments.

class Bound[source]

Bases: object

Type hint for make_funsor() decorated functions. This provides hints for bound variables (names).

class Has(bound)[source]

Bases: object

Type hint for make_funsor() decorated functions.

This hint asserts that a set of Bound variables always appear in the .inputs of the annotated argument.

For example, we could write a named matmul function that asserts that both arguments always contain the reduced input, and cannot be constant with respect to that input:

@make_funsor
def MatMul(
    x: Has[{"i"}],
    y: Has[{"i"}],
    i: Bound,
) -> Fresh[lambda x: x]:
    return (x * y).reduce(ops.add, i)

Here the string "i" in the annotations for x and y refer to the argument i of our MatMul function, which is known to be Bound (i.e it does not appear in the .inputs of evaluating Matmul(x, y, "i").

Warning

This annotation is experimental and may be removed in the future.

Note that because Funsor is inherently extensional, violating a Has constraint only raises a SyntaxWarning rather than a full TypeError and even then only under the reflect() interpretation.

As such, Has annotations should be used sparingly, reserved for cases where the programmer has complete control over the inputs to a function and knows that an argument will always depend on a bound variable, e.g. when writing one-off Funsor terms to describe custom layers in a neural network.

Parameters

bound (set) – A set of strings of names of Bound arguments of a make_funsor() -decorated function.

make_funsor(fn)[source]

Decorator to dynamically create a subclass of Funsor, together with a single default eager pattern.

This infers inputs, outputs, fresh, and bound variables from type hints follow the following convention:

  • Funsor inputs are typed Funsor.

  • Bound variable inputs (names) are typed Bound.

  • Fresh variable inputs (names) are typed Fresh together with lambda to compute the dependent domain.

  • Ground value inputs (e.g. Python ints) are typed Value together with their actual data type, e.g. Value[int].

  • The return value is typed Fresh together with a lambda to compute the dependent return domain.

For example to unflatten a single coordinate into a pair of coordinates we could define:

@make_funsor
def Unflatten(
    x: Funsor,
    i: Bound,
    i_over_2: Fresh[lambda i: Bint[i.size // 2]],
    i_mod_2: Fresh[lambda: Bint[2]],
) -> Fresh[lambda x: x]:
    assert i.output.size % 2 == 0
    return x(**{i.name: i_over_2 * Number(2, 3) + i_mod_2})
Parameters

fn (callable) – A type annotated function of Funsors.

Return type

subclas of Funsor