Is it possible to add a mode/separate interfaces that takes torch tensor and returns torch tensor

[larryliu0820]

Hi, I'm working on using torch_np to connect numpy to torch dynamo. One issue we are seeing is that all of the torch_np APIs are returning torch_np.ndarray. This means we have torch_np.ndarray in our captured fx graph. This is giving our downstream backends a hard time because they don't recognize torch_np.ndarray (and maybe shouldn't).

So I'm wondering if we can have a "mode" type of flip switch, to allow all the torch_np APIs to return tensor instead of torch_np.ndarray. Is this easy to do?

For an example of the captured graph:

def fn(x, y):
    a = x.numpy()
    b = y.numpy()
    return np.add(a, 1), np.add(b, 1) 

If we capture it through dynamo, it gives

[2023-05-15 11:40:25,582] torch._dynamo.output_graph.__graph: [DEBUG] TRACED GRAPH
 __compiled_fn_0 <eval_with_key>.0 opcode         name       target                               args             kwargs
-------------  ---------  -----------------------------------  ---------------  --------
placeholder    l_x_       L_x_                                 ()               {}
placeholder    l_y_       L_y_                                 ()               {}
call_function  ndarray    <class 'torch_np._ndarray.ndarray'>  (l_x_,)          {}
call_function  ndarray_1  <class 'torch_np._ndarray.ndarray'>  (l_y_,)          {}
call_function  add        <function add at 0x12de77a60>        (ndarray, 1)     {}
call_function  add_1      <function add at 0x12de77a60>        (ndarray_1, 1)   {}
output         output     output                               ((add, add_1),)  {}

Notice both a and b are being converted to torch_np.ndarray and their result after add is still an torch_np.ndarray. As a result, the outputs of the graph are also torch_np.ndarray. The backends to torch dynamo can't work with these graphs.

[larryliu0820]

cc @ezyang

[ezyang]

Our alternative is to metaprogram a big alternate API which is essentially torch_np as is, but with unwrapping on the outputs. Which we can do, it's just annoying.

[ev-br]

How does this differ from #1

[rgommers]

It seems to me that this is indeed gh-1, and that API-wise torch_np.add can only return torch_np.ndarray in order for things to work. What I am missing here is why the tracing cannot further unwrap the calls, why is that @larryliu0820 @ezyang? The call_function add that is shown in the graph is torch_np.add I assume - but in the end the actual addition that is being executed is torch.add which operates only on tensors, not wrapper ndarrays.

I admit not knowing much much about TorchDynamo internals, but we touched on this a couple of times and I also remember @lezcano saying "just keep tracing and it's all pytorch ops on tensors".

but with unwrapping on the outputs

I don't think this can be done per function in the torch_np namespace. If you have some switch that flips the output type of torch_np.add, then you'd change the meaning of something like torch_np.add(x.numpy(), 1).size, it'd now be a method (<function Tensor.size>) rather than a Python integer.

[lezcano]

+1 to Ralf. Why can't you unwrap these by accessing x.tensor on the outputs of the graph?

But at a more basic level, I think that TorchDynamo should know about torch_np.ndarrays, and it should use them to trace, but they should not appear in the trace. All the operations in torch_np are PyTorch and Python operations, so TorchDynamo should just register the PyTorch calls that appear in the call. Then, each torch_np.ndarray should have an associated Tensor variable, which is the one that will appear in the graph. Dynamo should then have a prologue and epilogue to this graph that wraps and unwraps NumPy tensors in these torch_np.ndarrays. The functions in torch_np cannot return Tensors themselves because of the reasons that Ralf described.

[ezyang]

The main problem is that we can't conveniently trace into torch_np code at Dynamo time; the most convenient thing to do is to trace torch_np operations directly into the produced Dynamo FX graph and rely on AOTAutograd tracing to eliminate them into ATen ops.

A more logical arrangement is that when Dynamo hits a numpy API call, it immediately lowers it into inner torch operations from torch_np so you never see torch_np in the Dynamo graph. However, this not something Dynamo really does today. In particular, we're not supposed to run AOTAutograd tracing until we're all done collecting the dynamo graph, so the only thing we really could do is have Dynamo inline the Python code from torch_np... which requires all of torch_np to be dynamo traceable. Which I suppose is not the worst thing in the world but it feels unnecessary because make_fx tracing should be fine for these ops.

Now, as for the other proposed changes:

What I am missing here is why the tracing cannot further unwrap the calls

We can, and @larryliu0820 has a prototype that adds another wrapper library around torch_np which does exactly this. The key distinction is that we have a Dynamo IR invariant, which is that all intermediates flowing through the graph are Tensor. So we cannot do the unwrapping in the IR, because that would imply you have a torch_np.ndarray temporarily before you unwrap it. We need a single function in the graph that does the wrapping-torch_np-call-unwrapping all in one go. We can generate this, and it seems like this is what we're going to do.

Dynamo should then have a prologue and epilogue to this graph that wraps and unwraps NumPy tensors in these torch_np.ndarrays.

This proposal violates the "only Tensor intermediates" invariant on the FX graph. It also is a bit difficult to implement in practice, because the epilogue code for unwrapping must happen INSIDE the graph, because an even more important invariant is "backends only need to compile functions that take/produce Tensors (not torch_np.ndarray)"