Full-featured distribution wrappers

[eb8680]

The biggest single issue blocking wider use of Funsor in Pyro and NumPyro right now is the incomplete coverage of distributions.

At a high level, the goal is to be able to perform all distribution operations that appear in any Pyro or NumPyro model (e.g. sampling and scoring) on Funsors directly, where distribution funsors are obtained from using to_funsor to automatically convert distributions to Funsors initially and using funsor.to_data to automatically convert the final results back to raw PyTorch/JAX objects. Wherever possible, the Funsor wrappers should also avoid the need for user-facing higher-order distributions, such as Independent or TransformedDistribution, in favor of idiomatic Funsor operations or broadcasting semantics.

We have gotten pretty far with the generic wrappers in funsor.distribution and funsor.{jax,torch}.distributions, but finishing the job and achieving full coverage of pyro.distributions remains a challenge because of the large distribution API and number of distributions, many small impedance mismatches and legacy design choices in PyTorch (e.g. the data type of Bernoulli), and difficulty of programmatic access and automation (e.g. no generic tool for constructing random valid instances of a distribution given a batch_shape).

I've tried to collect the remaining Funsor-specific tasks in this issue so that we can better measure progress toward this goal. We may also need to do additional work upstream in Pyro, NumPyro or PyTorch distributions.

Transforms and TransformedDistributions (some design discussion in #309):

• Add TanhOp and AtanhOp #387 Near-complete coverage of parameter-free invertible Transforms in Pyro
• Support conversion of Transforms and TransformedDistributions to and from Funsors #365 Basic conversion of TransformedDistributions to and from funsor.Distributions on the PyTorch backend
• Add factory for dynamic TransformOp #427 Automatic wrapping of non-batched Transforms
• Conversion of TransformedDistributions to and from funsors on the JAX backend (can copy Support conversion of Transforms and TransformedDistributions to and from Funsors #365)
• Coverage for distribution conversion with non-invertible Transforms like AbsTransform
• Coverage for conversion of invertible Transforms with ground parameters, notably AffineTransform and PowerTransform with batch_shape == ()
• Coverage for conversion of invertible Transform with non-ground parameters, e.g. AffineTransform with len(batch_shape) > 0
• Wrappers for ConditionalTransforms and TransformedModules in Pyro/PyTorch

Other basic distribution modifiers:

• Tweak automated registration for distribution-to-funsor conversion #391 automated to_funsor conversion for custom distributions
• Add support for most remaining atomic distributions #394 Near-full coverage of basic distributions (easy enough to implement; the difficulty is testing and fixing PyTorch bugs)
• Add support for most remaining atomic distributions #394 Get backend Delta conversion working
• Support converting funsor.terms.Independent to and from data #396 Converting funsor.Independent to and from Independent distributions
• Update distribution shape inference to handle independent dims #402 more idiomatic support for Independent distributions in Funsor
• Add method for converting Pyro ExpandedDistribution to funsor #418 Converting lazy Pyro ExpandedDistributions to funsors
• support ExpandedDistribution for jax backend #419 Converting lazy Numpyro ExpandedDistributions to funsors
• Support IndependentConstraint in Distribution._infer_param_domain #432 Handle conversion of parameters with _IndependentConstraint arg_constraints
• Directly convert IndependentDistribution to base distribution #443 Convert IndependentDistributions directly to base Funsor distributions
• Support for converting mixture distributions into funsor sum-product expressions, e.g. MaskedMixture and MixtureSameFamily

Masking:

• First-class support for masking using logical operations rather than floats
• Optimized handling for scalar masks (i.e. the case mask is False) (discussed in Support mask=False for backend MaskedDistribution #459)

Direct TFP distribution wrappers:

• Direct wrapping of TFP distributions in the JAX backend - this would probably involve a new subclass class TFPDistribution(funsor.distribution.Distribution) with TFP-specific implementations of _infer_value_domain and _infer_param_domain
• Direct TFP Bijector wrappers

Atomic distribution computations beyond sampling and scoring implemented in the backend libraries:

• Add generic stat methods to Distribution #388 Entropy (discussed in Exploit opportunities for analytic KL and entropy computations #374)
• Add generic stat methods to Distribution #388 Mean (TorchDistribution.mean)
• Add generic stat methods to Distribution #388 Variance (TorchDistribution.variance)
• KL divergence (discussed in Exploit opportunities for analytic KL and entropy computations #374)

Test harnesses for distribution wrappers (testing correctness of underlying distribution functionality here is out of scope - we are mostly interested in ensuring that results are converted to Funsors correctly):

• Automate distribution testing #389 conversion with to_funsor and to_data
• Automate distribution testing #389 types and shapes
• Automate distribution testing #389 densities
• Automate distribution testing #389 enumeration (enumerate_support)
• Automate distribution testing #389 samplers
• transforms

Miscellaneous:

• Find a workaround for int/float casting issues in Bernoulli (discussed in Casting between real and bint #348)
• Generic support for sampling via Funsor.sample() in funsor.pyro.FunsorDistribution - done properly, this may eliminate the need for first-class conjugate distribution implementations in backends, e.g. DirichletMultinomial

Lower priority, possibly unnecessary:

• conjugate_update wrapper (may be unnecessary?)
• Isolate Monte Carlo distribution sample and gradient tests on a new branch #395 Monte Carlo tests for samplers and their gradients (low priority)
• Some Funsor analogue of Distribution.expand() e.g. via lazy units as in Add lazy multidimensional Unit terms #235 or substitution of indexed Variable terms (f(v=v['i']))
• Generic tests for conjugate distribution pairs
• deprecate existing, case-by-case tests

[eb8680]

Conversion of TransformedDistributions to and from funsors on the JAX backend (can copy #365)

Note this will be a bit trickier than expected because in NumPyro Transform.inv is just a regular method, rather than returning an _InverseTransform, which NumPyro does not implement.

[fehiepsi]

I will make a PR for that. :)