refactor Gaussian to use Tensor api instead of torch.Tensor

[fehiepsi]

This PR follows Approach 2 in #207 to make Gaussian backend-agnostic. My plan is to

• check if info_vec is torch.Tensor, then set self.backend = TorchTensor.
• make backend-specific ops such as triangular_solve, cholesky static methods of their corresponding TorchTensor type. For example,

x.triangular_solve(y)

will be

self.backend.triangular_solve(x, y)

It seems that there is a lot of work to do but I hope that future refactoring will be easier after this.

Tasks

• [] subclass Tensor to TorchTensor this is not necessary for now
• collect and implement all required static methods ops for TorchTensor
• use the above static methods ops in Gaussian
• do the same for other classes/functions in gaussian.py
• make sure all tests pass
• make align_tensors and materialize backend agnostic
• make numpy backend pass for gaussian smoke test

Future PRs

• deal with torch._C._get_tracing_state() statements
• make sure that numpy backend will work

[fritzo]

"... methods of their corresponding TorchTensor type"

Instead of adding methods, let's add ops, following e.g. matmul or sigmoid or pow

# in ops.py
class TriangularSolveOp(Op):
    pass

@TriangularSolveOp
def triangular_solve(x, y):
    raise NotImplementedError

# in torch.py
@eager.register(Binary, TriangularSolveOp, Tensor, Tensor)
def triangular_solve(x, y):
    assert len(x.output.shape) == 2
    assert len(y.output.shape) == 1  # is this right?
    ...other assertions...
    inputs, x, y = align_tensors(x, y)
    data = x.triangular_solve(y)
    return Tensor(data, inputs)

# in numpy.py
@eager.register(Binary, TriangularSolveOp, Array, Array)
def triangular_solve(x, y):
    ...

Note it's probably cleanest to add all necessary ops in a separate PR.

[fehiepsi]

Thank you, @fritzo! I am not familiar with dispatching register so based on your suggestion, I tried to port torch.cholesky and torch.cat to ops. And it seems to work out of the box. :D I think that other ops can be implemented similarly (for ops.zeros, maybe we can have ops.new_zeros(prototype, shape)). Right now, I am just exploring the functionalities of funsor. When I understand it more, I'll make a separate PR with necessary ops for clarity.

It seems that your suggestion based on the assumption that info_vec and precision are Tensor/Array. Do you suggest to have Gaussian(Tensor, Tensor) instead of Gaussian(torch.Tensor, torch.Tensor, inputs)? In addition, when do we want to use @eager.register(Binary, TriangularSolveOp, Tensor, Tensor)? It seems to me that it is easier to implement @triangular_solve.register(torch.Tensor, torch.Tensor).

[fritzo]

Do you suggest to have Gaussian(Tensor, Tensor) instead of Gaussian(torch.Tensor, torch.Tensor, inputs)?

Well @eb8680 and I had originally envisioned supporting Gaussian(Funsor,Funsor,inputs), probably with an assertion for isinstance(_, (Tensor, Array)) as in Delta. However this approach is more extreme than your original approach of allowing multiple backends, and it would involve lots of subtle issues. I guess it is best to pursue the Gaussian(torch.Tensor,torch.Tensor,inputs) approach initially.

One issue blocking the Gaussian(Funsor,Funsor) approach is that we will want to force internal ops to be naive-eagerly evaluated; if we simply replace Gaussian's PyTorch guts with pure funsor, then we'll need to revise our interpretation ordering. Currently the eager and normalize interpretations are interleaved in a way that makes arithmetic expensive. In the Gaussian guts we have hand-optimized expressions to already be normalized and want a pure eager evaluation. This is a deep issue beyond the Gaussian implementation, but the Gaussian implementation is a great motivating example to make our interpretations more flexible. I think what we'll really want is for normalize to be an interpretation->interpreation transform rather than an interpretation, then we can use with interpretation(normalize(eager))) in user code and with interpreatation(eager) in hand-normalized library code. This is a real research question!

[fritzo]

In addition, when do we want to use @eager.register(Binary, TriangularSolveOp, Tensor, Tensor)? It seems to me that it is easier to implement @triangular_solve.register(torch.Tensor, torch.Tensor).

Good point. Probably it is easiest to start with the latter, and make Gaussian accept either torch.Tensor or numpy.ndrray inputs.

[eb8680]

In the Gaussian guts we have hand-optimized expressions to already be normalized and want a pure eager evaluation

I think the easy way to do this is just to write a separate eager_only interpretation and decorate any functions whose internals we want to be eagerly and opaquely evaluated:

def eager_only(...):
    result = eager.dispatch(...)
    if result is None:
        result = reflect(...)
    return result

@eager.register(Binary, AddOp, Gaussian, Gaussian)
@interpretation(eager_only)
def eager_binary_gaussian_gaussian(op, lhs, rhs):
    ...

Currently the eager and normalize interpretations are interleaved in a way that makes arithmetic expensive

The above is a useful fix for this issue, but more generally note that if an expression can be eagerly evaluated as-is then normalize will never be called.

[fehiepsi]

@fritzo This should be good to go for now. I'll deal with torch._C._get_tracing_state() statements and supports for numpy backend in another PR. In this PR, all torch.Tensor ops are deferred to funsor.ops.

[fritzo]

It would be nice to either update or fork tests.test_gaussian.test_fork() to ensure all ops work with numpy.ndarrays also. What do you think of adding a backend param like this?

@pytest.mark.parametrize('expr,expected_type', [
    ('-g1', Gaussian),
    ...
    ('g1(x=x0)', (Tensor, Array)),  # allow either backend
    ...
])
@pytest.mark.parametrize("backend", ["torch", "numpy"])
def test_smoke(backend, expr, expected_type):
    if backend == "torch":
        tensor = torch.tensor
        Tensor = Tensor
    else:
        tensor = np.array
        Tensor = Array

    ...continue with rest of test...

EDIT Or did you have another testing strategy in mind?

[fritzo]

I think you'll need to add scipy to the test requirements in setup.py

[fehiepsi]

Yes, I'll do that. I think numpy backend will work after I make align_tensor and materialize backend agnostic.

[fehiepsi]

@fritzo Smoke tests are passed now. For further gaussian tests, I'll address issues with numpy backend in another PR, which might require changes in other files beside gaussian.py (I'm not sure). In last commits, I added ops.Tensor, ops.align_tensor_op, and ops.materialize with some changes:

• add expand arg to numpy align_array
• add prototype arg to materialize. The reason is if the input is not Number/Array/Tensor, then we don't know if we need to use torch.arange or numpy.arange. So prototype is helpful to distinguish them.
• fix the issue of numpy triangular solve does not support batching.

[fritzo]

This looks mostly good, but we'll need to change align_tensors and materialize.

Funsor is organized at varying levels of abstraction. The lowest level of abstraction is ops which represents pointwise mathematical operations. Higher level than ops are: domains and then terms including free variables and algebra involving ops. You can think of ops as arithmetic and terms as algebra.

While most of the ops you've added are indeed arithmetic operations, the two functions align_tensor and materialize involve algebraic concepts of inputs and free variables. To maintain our abstraction hierarchy, I think a cleaner generic design would be to add @singledispatch or @multipledispatch wrappers of these functions in gaussian.py. Can you find a way to maintain this hieararchy?

[fehiepsi]

Thanks for explaining, @fritzo ! I'll find solutions for them. It should be doable. :)

[fehiepsi]

@fritzo I just use dispatch for a hidden _materialize and add permute, new_arange ops so that align_tensor is backend agnostic.

[fritzo]

These look great! Do you think we should them into say a funsor/tensor_ops.py, delete the original versions in funsor/torch.py and funsor/numpy.py? That way torch.py, numpy.py and gaussian.py could all import these two generic functions.

[fehiepsi]

To me, we should do so. What is your opinion?

[fritzo]

Let's do it. I'm happy merging as is and deferring to a follow-up PR. Or making the change in this PR. Whichever you prefer.

[fritzo]

Let's just move that to a follow-up PR. This PR is already getting big 🙂

[fehiepsi]

Agreed! :)