Support multiple backends

[fritzo]

This issue tracks changes needed to fully support multiple backends in an interchangeable/generic way. The initial release of Funsor implements most funsors using PyTorch (funsor.torch implements Tensor, but all of funsor.gaussian, .joint, .distributions use PyTorch internally.

Approach 1

One possibility is to replicate backend-specific logic in a consistent way, e.g. with module structure

• funsor
  • torch
    • tensor
    • gaussian
    • cnf
    • distributions
  • numpy / jax
    • tensor
    • gaussian
    • cnf
    • distributions
  • tensorflow
    • ...
  • birch
    • ...

Approach 2

A second possibility (suggested by @eb8680 below and in #159) is to

1. subclass Tensor into TorchTensor and NumpyTensor, then
2. refactor Gaussian and Contraction etc. to use Tensors rather than torch.Tensors internally.

This approach would still require duplicating distributions.

Potential backends

• PyTorch / pyro.distributions
• Add numpy backend for funsor #58 NumPy
• JAX / numpyro.distributions
• Tensorflow 2.0 / tfp.distributions
• Trax's TensorFlow-backed NumPy API tf_numpy
• Birch

Other tasks

• Remove stray dependencies on PyTorch or NumPy e.g. in funsor.testing and funsor.domains
• Add backend-specific installation options funsor[torch], funsor[numpy] etc.

[eb8680]

Related suggestion in #159. Agree we should be able to make Number, Tensor, Gaussian, Delta, Distribution interfaces backend-agnostic and have backend independence follow easily from there.

[eb8680]

I believe the bulk of this work is finished for the PyTorch and JAX backends as of #327, but I'll leave this issue open until #332 is addressed.

[eb8680]

I'm going to close this as resolved. If we want to add more backends in the future or make Gaussian use funsor.Tensor for its parameters, we can track progress in separate issues.