Don't replace unused variables with None

[beverlylytle]

Created, but unused variables in a Thunder jitted program are replaced by a placeholder None, which, then in the executable is replaced by _, as is conventional in Python. More concretely,

def f(x):
    y, z = torch.var_mean(x)
    return y

would be represented by a trace with a bound symbol for the var_mean op with output consisting of a TensorProxy and None, which then would be executed as something like y, _ = torch.var_mean(x). Not only is it confusing to see None seemingly being assigned a value, this variable _ takes up memory and is never deleted. This PR leaves the unused variable as a TensorProxy so that it can be removed with del_last_used.

[beverlylytle]

Some grad tests are failing in a non-trivial way. In the case of atleast_2d, test_grad is being supplied an input which is then directly returned as output (z, y = atleast_2d(x, y)). Before this change the output y would have been replaced by None and we wouldn't hit the error raised here.

[IvanYashchuk]

Could you please check whether the error goes away if atleast_2d and all similar operations use prims.clone and not return the input directly?

[t-vi]

I would not expect clone (which is not what PyTorch does, it really does identity.) Maybe what is missing is something in __iter__.
Note that the grad checks are still using the old logic, which has much more elaborate handling of complex outputs, so maybe moving those from running vjp to comparing to torch grad after torch.autograd.grad(thunder_function(x), x) fixes it. I would recommend looking into whether it is test specific before trying to fix. I vaguely recall a similar problem before but I don't remember the exact issue.

[beverlylytle]

You are right, @t-vi , that prims.clone is not the right approach here because it does not mimic PyTorch. However I don't think this is just limited to the tests. For example,

import torch
import thunder

class MyFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x: torch.Tensor) -> torch.Tensor:
        x = x.sin()
        return torch.atleast_2d(x)

    @staticmethod
    def backward(ctx, grad_output) -> torch.Tensor:
        return grad_output

class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x) -> torch.Tensor:
        return MyFunction.apply(x)

model = Model().to(dtype=torch.float64)

jitted = thunder.jit(model)

y = torch.randn(10, 10, requires_grad=True, device="cuda")
z = jitted(y)

still fails with "Encountered exception ValueError: Variable t1 is being overwritten this is not allowed while tracing Model()", though this failure is unrelated to removing this helper function. Using prims.shallow_copy fixes both issues.

[beverlylytle]

Btw, which __iter__ do you mean @t-vi ?

[t-vi]

Couldn't we del _ instead when it is assigned?

[beverlylytle]

import torch
import thunder

def f(x):
    y, z = torch.var_mean(x, dim=1)
    return z

jitted = thunder.jit(f)

x = torch.randn(10, 10)
z = jitted(x)

last_trace = thunder.last_traces(jitted)[-1]
var_mean_bsym = last_trace.bound_symbols[1]
print(var_mean_bsym.flat_outs)         # [None, <TensorProxy(name="t11", dtype=thunder.dtypes.float32, shape=(10,))>]

print(var_mean_bsym)                   # (_, t11) = torch.var_mean(x, 1, keepdim=False, correction=1)
                                       #   # (t10, t11) = ltorch.var_mean(x, 1, keepdim=False, correction=1)
                                       #     # (t10, t11) = prims.var_mean(x, (1,), correction=1)

The underscore doesn't appear until the printing process and is otherwise represented as None. I attempted to modify del_last_used to also add a bsym with op prims.DEL and arg None, but this meant also changing the del_printer from arg_string = ", ".join(codeutils.prettyprint(x, literals_allowed=False) for x in arg_printables) to arg_string = ", ".join(codeutils.prettyprint(x, literals_as_underscores=True) for x in arg_printables). But that gives me the howling fantods.

[t-vi]

I think the del_last_used could check if the return is a tuple and if it is it can just iterate over that instead of using flat proxy outs. This would allow you to detect none properly.

[beverlylytle]

@t-vi I'm not understanding your last comment, or at least the motivation behind it. Correct me if I'm wrong, but it seems like you are suggesting that unused TensorProxies should continue to be replaced with None in the dce pass and that del_last_used should then tidy things up. What is the advantage of replacing an unused TensorProxy with None in a bsym's outs?

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[t-vi]

Thank you @beverlylytle @IvanYashchuk @crcrpar