Funsors¶
Basic Funsors¶
-
reflect
(cls, *args, **kwargs)[source]¶ Construct a funsor, populate
._ast_values
, and cons hash. This is the only interpretation allowed to construct funsors.
-
eager_or_die
(cls, *args)[source]¶ Eagerly execute ops with known implementations. Disallows lazy
Subs
,Unary
,Binary
, andReduce
.Raises: NotImplementedError
no pattern is found.
-
sequential
(cls, *args)[source]¶ Eagerly execute ops with known implementations; additonally execute vectorized ops sequentially if no known vectorized implementation exists.
-
moment_matching
(cls, *args)[source]¶ A moment matching interpretation of
Reduce
expressions. This falls back toeager
in other cases.
-
class
Funsor
(inputs, output, fresh=None, bound=None)[source]¶ Bases:
object
Abstract base class for immutable functional tensors.
Concrete derived classes must implement
__init__()
methods taking hashable*args
and no optional**kwargs
so as to support cons hashing.Derived classes with
.fresh
variables must implement aneager_subs()
method. Derived classes with.bound
variables must implement an_alpha_convert()
method.Parameters: - inputs (OrderedDict) – A mapping from input name to domain. This can be viewed as a typed context or a mapping from free variables to domains.
- output (Domain) – An output domain.
-
dtype
¶
-
shape
¶
-
requires_grad
¶
-
sample
(sampled_vars, sample_inputs=None)[source]¶ Create a Monte Carlo approximation to this funsor by replacing functions of
sampled_vars
withDelta
s.The result is a
Funsor
with the same.inputs
and.output
as the original funsor (plussample_inputs
if provided), so that self can be replaced by the sample in expectation computations:y = x.sample(sampled_vars) assert y.inputs == x.inputs assert y.output == x.output exact = (x.exp() * integrand).reduce(ops.add) approx = (y.exp() * integrand).reduce(ops.add)
If
sample_inputs
is provided, this creates a batch of samples scaled samples.Parameters:
-
unscaled_sample
(sampled_vars, sample_inputs)[source]¶ Internal method to draw an unscaled sample. This should be overridden by subclasses.
-
align
(names)[source]¶ Align this funsor to match given
names
. This is mainly useful in preparation for extracting.data
of afunsor.torch.Tensor
.Parameters: names (tuple) – A tuple of strings representing all names but in a new order. Returns: A permuted funsor equivalent to self. Return type: Funsor
-
eager_subs
(subs)[source]¶ Internal substitution function. This relies on the user-facing
__call__()
method to coerce non-Funsors to Funsors. Once all inputs are Funsors,eager_subs()
implementations can recurse to callSubs
.
-
to_data
(x)[source]¶ Extract a python object from a
Funsor
.Raises a
ValueError
if free variables remain or if the funsor is lazy.Parameters: x – An object, possibly a Funsor
.Returns: A non-funsor equivalent to x
.Raises: ValueError if any free variables remain. Raises: PatternMissingError if funsor is not fully evaluated.
-
class
Variable
(name, output)[source]¶ Bases:
funsor.terms.Funsor
Funsor representing a single free variable.
Parameters: - name (str) – A variable name.
- output (funsor.domains.Domain) – A domain.
-
class
Subs
(arg, subs)[source]¶ Bases:
funsor.terms.Funsor
Lazy substitution of the form
x(u=y, v=z)
.Parameters: - arg (Funsor) – A funsor being substituted into.
- subs (tuple) – A tuple of
(name, value)
pairs, wherename
is a string andvalue
can be coerced to aFunsor
viato_funsor()
.
-
class
Unary
(op, arg)[source]¶ Bases:
funsor.terms.Funsor
Lazy unary operation.
Parameters:
-
class
Binary
(op, lhs, rhs)[source]¶ Bases:
funsor.terms.Funsor
Lazy binary operation.
Parameters:
-
class
Reduce
(op, arg, reduced_vars)[source]¶ Bases:
funsor.terms.Funsor
Lazy reduction over multiple variables.
Parameters:
-
class
Number
(data, dtype=None)[source]¶ Bases:
funsor.terms.Funsor
Funsor backed by a Python number.
Parameters: - data (numbers.Number) – A python number.
- dtype – A nonnegative integer or the string “real”.
-
class
Slice
(name, start, stop, step, dtype)[source]¶ Bases:
funsor.terms.Funsor
Symbolic representation of a Python
slice
object.Parameters:
-
class
Stack
(name, parts)[source]¶ Bases:
funsor.terms.Funsor
Stack of funsors along a new input dimension.
Parameters:
-
class
Cat
(name, parts, part_name=None)[source]¶ Bases:
funsor.terms.Funsor
Concatenate funsors along an existing input dimension.
Parameters:
-
class
Lambda
(var, expr)[source]¶ Bases:
funsor.terms.Funsor
Lazy inverse to
ops.getitem
.This is useful to simulate higher-order functions of integers by representing those functions as arrays.
Parameters: - var (Variable) – A variable to bind.
- expr (funsor) – A funsor.
-
class
Independent
(fn, reals_var, bint_var, diag_var)[source]¶ Bases:
funsor.terms.Funsor
Creates an independent diagonal distribution.
This is equivalent to substitution followed by reduction:
f = ... # a batched distribution assert f.inputs['x_i'] == reals(4, 5) assert f.inputs['i'] == bint(3) g = Independent(f, 'x', 'i', 'x_i') assert g.inputs['x'] == reals(3, 4, 5) assert 'x_i' not in g.inputs assert 'i' not in g.inputs x = Variable('x', reals(3, 4, 5)) g == f(x_i=x['i']).reduce(ops.logaddexp, 'i')
Parameters:
-
to_funsor
(*args, **kwargs)¶ Multiply dispatched method: to_funsor
- Convert to a
Funsor
. Only
Funsor
s and scalars are accepted.param x: An object. param funsor.domains.Domain output: An optional output hint. return: A Funsor equivalent to x
.rtype: Funsor raises: ValueError - Other signatures:
- Funsor object, object object, Domain Funsor, Domain str, Domain Number Number, Domain Tensor Tensor, Domain ndarray ndarray, Domain
- Convert to a
Delta¶
PyTorch¶
-
align_tensor
(new_inputs, x, expand=False)[source]¶ Permute and add dims to a tensor to match desired
new_inputs
.Parameters: - new_inputs (OrderedDict) – A target set of inputs.
- x (funsor.terms.Funsor) – A
Tensor
orNumber
. - expand (bool) – If False (default), set result size to 1 for any input
of
x
not innew_inputs
; if True expand tonew_inputs
size.
Returns: a number or
torch.Tensor
that can be broadcast to other tensors with inputsnew_inputs
.Return type: int or float or torch.Tensor
-
align_tensors
(*args, **kwargs)[source]¶ Permute multiple tensors before applying a broadcasted op.
This is mainly useful for implementing eager funsor operations.
Parameters: - *args (funsor.terms.Funsor) – Multiple
Tensor
s andNumber
s. - expand (bool) – Whether to expand input tensors. Defaults to False.
Returns: a pair
(inputs, tensors)
where tensors are alltorch.Tensor
s that can be broadcast together to a single data with giveninputs
.Return type: - *args (funsor.terms.Funsor) – Multiple
-
class
Tensor
(data, inputs=None, dtype='real')[source]¶ Bases:
funsor.terms.Funsor
Funsor backed by a PyTorch Tensor.
This follows the
torch.distributions
convention of arranging named “batch” dimensions on the left and remaining “event” dimensions on the right. The output shape is determined by all remaining dims. For example:data = torch.zeros(5,4,3,2) x = Tensor(data, OrderedDict([("i", bint(5)), ("j", bint(4))])) assert x.output == reals(3, 2)
Operators like
matmul
and.sum()
operate only on the output shape, and will not change the named inputs.Parameters: - data (torch.Tensor) – A PyTorch tensor.
- inputs (OrderedDict) – An optional mapping from input name (str) to
datatype (
Domain
). Defaults to empty. - dtype (int or the string "real".) – optional output datatype. Defaults to “real”.
-
requires_grad
¶
-
arange
(name, *args, **kwargs)[source]¶ Helper to create a named
torch.arange()
funsor. In some cases this can be replaced by a symbolicSlice
.Parameters: Return type:
-
materialize
(x)[source]¶ Attempt to convert a Funsor to a
Number
orTensor
by substitutingarange()
s into its free variables.Parameters: x (Funsor) – A funsor. Return type: Funsor
-
class
Function
(fn, output, args)[source]¶ Bases:
funsor.terms.Funsor
Funsor wrapped by a PyTorch function.
Functions are assumed to support broadcasting and can be eagerly evaluated on funsors with free variables of int type (i.e. batch dimensions).
Function
s are usually created via thefunction()
decorator.Parameters: - fn (callable) – A PyTorch function to wrap.
- output (funsor.domains.Domain) – An output domain.
- args (Funsor) – Funsor arguments.
-
function
(*signature)[source]¶ Decorator to wrap a PyTorch function.
Example:
@funsor.torch.function(reals(3,4), reals(4,5), reals(3,5)) def matmul(x, y): return torch.matmul(x, y) @funsor.torch.function(reals(10), reals(10, 10), reals()) def mvn_log_prob(loc, scale_tril, x): d = torch.distributions.MultivariateNormal(loc, scale_tril) return d.log_prob(x)
To support functions that output nested tuples of tensors, specify a nested tuple of output types, for example:
@funsor.torch.function(reals(8), (reals(), bint(8))) def max_and_argmax(x): return torch.max(x, dim=-1)
Parameters: *signature – A sequence if input domains followed by a final output domain or nested tuple of output domains.
-
class
Einsum
(equation, operands)[source]¶ Bases:
funsor.terms.Funsor
Wrapper around
torch.einsum()
to operate on real-valued Funsors.Note this operates only on the
output
tensor. To perform sum-product contractions on named dimensions, instead use+
andReduce
.Parameters: - equation (str) – An
torch.einsum()
equation. - operands (tuple) – A tuple of input funsors.
- equation (str) – An
-
torch_tensordot
(x, y, dims)[source]¶ Wrapper around
torch.tensordot()
to operate on real-valued Funsors.Note this operates only on the
output
tensor. To perform sum-product contractions on named dimensions, instead use+
andReduce
.Arguments should satisfy:
len(x.shape) >= dims len(y.shape) >= dims dims == 0 or x.shape[-dims:] == y.shape[:dims]
Parameters: Return type:
NumPy¶
-
align_array
(new_inputs, x)[source]¶ Permute and expand an array to match desired
new_inputs
.Parameters: - new_inputs (OrderedDict) – A target set of inputs.
- x (funsor.terms.Funsor) – A
Array
s or orNumber
.
Returns: a number or
numpy.ndarray
that can be broadcast to other array with inputsnew_inputs
.Return type:
-
align_arrays
(*args)[source]¶ Permute multiple arrays before applying a broadcasted op.
This is mainly useful for implementing eager funsor operations.
Parameters: *args (funsor.terms.Funsor) – Multiple Array
s andNumber
s.Returns: a pair (inputs, arrays)
where arrayss are allnumpy.ndarray
s that can be broadcast together to a single data with giveninputs
.Return type: tuple
-
class
ArrayMeta
(name, bases, dct)[source]¶ Bases:
funsor.terms.FunsorMeta
Wrapper to fill in default args and convert between OrderedDict and tuple.
-
class
Array
(data, inputs=None, dtype='real')[source]¶ Bases:
funsor.terms.Funsor
Funsor backed by a NumPy Array.
This follows the
torch.distributions
convention of arranging named “batch” dimensions on the left and remaining “event” dimensions on the right. The output shape is determined by all remaining dims. For example:data = np.zeros((5,4,3,2)) x = Array(data, OrderedDict([("i", bint(5)), ("j", bint(4))])) assert x.output == reals(3, 2)
Operators like
matmul
and.sum()
operate only on the output shape, and will not change the named inputs.Parameters:
-
arange
(name, size)[source]¶ Helper to create a named
numpy.arange()
funsor.Parameters: Return type:
-
materialize
(x)[source]¶ Attempt to convert a Funsor to a
Number
ornumpy.ndarray
by substitutingarange()
s into its free variables.
Gaussian¶
-
class
BlockVector
(shape)[source]¶ Bases:
object
Jit-compatible helper to build blockwise vectors. Syntax is similar to
torch.zeros()
x = BlockVector((100, 20)) x[..., 0:4] = x1 x[..., 6:10] = x2 x = x.as_tensor() assert x.shape == (100, 20)
-
class
BlockMatrix
(shape)[source]¶ Bases:
object
Jit-compatible helper to build blockwise matrices. Syntax is similar to
torch.zeros()
x = BlockMatrix((100, 20, 20)) x[..., 0:4, 0:4] = x11 x[..., 0:4, 6:10] = x12 x[..., 6:10, 0:4] = x12.transpose(-1, -2) x[..., 6:10, 6:10] = x22 x = x.as_tensor() assert x.shape == (100, 20, 20)
-
align_gaussian
(new_inputs, old)[source]¶ Align data of a Gaussian distribution to a new
inputs
shape.
-
class
Gaussian
(info_vec, precision, inputs)[source]¶ Bases:
funsor.terms.Funsor
Funsor representing a batched joint Gaussian distribution as a log-density function.
Mathematically, a Gaussian represents the density function:
f(x) = < x | info_vec > - 0.5 * < x | precision | x > = < x | info_vec - 0.5 * precision @ x >
Note that
Gaussian
s are not normalized, rather they are canonicalized to evaluate to zero log density at the origin:f(0) = 0
. This canonical form is useful in combination with the information filter representation because it allowsGaussian
s with incomplete information, i.e. zero eigenvalues in the precision matrix. These incomplete distributions arise when making low-dimensional observations on higher dimensional hidden state.Parameters: - info_vec (torch.Tensor) – An optional batched information vector,
where
info_vec = precision @ mean
. - precision (torch.Tensor) – A batched positive semidefinite precision matrix.
- inputs (OrderedDict) – Mapping from name to
Domain
.
- info_vec (torch.Tensor) – An optional batched information vector,
where
Joint¶
Contraction¶
-
class
Contraction
(red_op, bin_op, reduced_vars, terms)[source]¶ Bases:
funsor.terms.Funsor
Declarative representation of a finitary sum-product operation.
After normalization via the
normalize()
interpretation contractions will canonically order their terms by type:Delta, Number, Tensor, Gaussian
-
GaussianMixture
¶ alias of
funsor.cnf.Contraction
Integrate¶
-
class
Integrate
(log_measure, integrand, reduced_vars)[source]¶ Bases:
funsor.terms.Funsor
Funsor representing an integral wrt a log density funsor.
Parameters: