Skip to content

Add Funsor (Constant) that is constant wrt to input variables #547

Description

@ordabayevy

Sometimes it is useful to keep track of certain Funsor inputs even though the underlying Funsor is constant wrt to those inputs. Examples are Zero terms in Delta.eager_reduce and In contrib.funsor.TraceEnum_ELBO. This is currently implemented as expanded Tensor terms which is memory/computation inefficient and only allows Bint input type. Proposed Constant funsor has following benefits:

  1. Declare constant inputs of a Funsor in a memory efficient way: e.g. Constant(cost.inputs, Number(0))
  2. Allow variables with Real output type: e.g. Constant(OrderedDict(x=Bint[3], y=Real), Number(0))
  3. Define computationally efficient patterns for subs/reduce/unary/binary operations
  4. Convert to and from ProvenanceTensor (ProvenanceTensor #543) where provenance of the tensor corresponds to const_inputs. Edit: I don’t think this is correct, the counterpart of ProvenanceTensor should something similar to Delta.
# Substitution

# 0(x, y)
a = Constant(OrderedDict(x=Bint[3], y=Real), Number(0))
# 0(x=0, y) = 0(y)
assert a(x=0) is Constant(OrderedDict(y=Real), Number(0))
# 0(x=0, y=3.14) = 0
assert a(x=0, y=3.14) is Number(0)

# Binary ops

# 1(x, y)
a = Constant(OrderedDict(x=Bint[3], y=Real), Number(1))
# b(x)
b = Tensor(...)["x"]
# 1(x, y) *  b(x) = (1 * b(x))(y) = b(x)(y)
c = a * b  # returns Constant(OrderedDict(y=Real), b*Number(1))

# Reduction

# 1(x, y)
a = Constant(OrderedDict(x=Bint[3], y=Real), Number(1))
# 1(x=0, y) + 1(x=1, y) + 1(x=2, y) = 3(y)
a.reduce(ops.add, "x")  # returns Constant(OrderedDict(y=Real), Number(3))

The argument in Constant(const_inputs, arg) doesn't have to be a Number. It can be any Funsor with the condition that const_inputs and arg.inputs have to be disjoint.

The main motivation is to use Constant as a wrapper for ProvenanceTensor and as targets for log_measures in a general version of contrib.fusor.Trace_ELBO (pyro-ppl/pyro#2893).

Metadata

Metadata

Assignees

No one assigned

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions