fx: unable to symbolically trace model with torch.zeroes(Proxy)

[vkuzo]

Currently using proxied tensor shapes as inputs to torch.zeros is not symbolically traceable:

# simple example

import torch

class M(torch.nn.Module):
    def forward(self, x):
        x = torch.zeros(x.shape[0], x.shape[1])
        return x
    
m = M()
gm = torch.fx.symbolic_trace(m)
print(gm)

...

<ipython-input-9-de0156858ee1> in forward(self, x)
      3 class M(torch.nn.Module):
      4     def forward(self, x):
----> 5         x = torch.zeros(x.shape[0], x.shape[1])
      6         return x
      7 

TypeError: zeros() received an invalid combination of arguments - got (Proxy, Proxy), but expected one of:
 * (tuple of ints size, *, tuple of names names, torch.dtype dtype, torch.layout layout, torch.device device, bool pin_memory, bool requires_grad)
 * (tuple of ints size, *, Tensor out, torch.dtype dtype, torch.layout layout, torch.device device, bool pin_memory, bool requires_grad)

# real model code

    def forward(self, x1, x2):
        assert x1.shape == x2.shape
        d = self.d
        s1 = self.s1
        s2 = self.s2
        n, c, h, w = x1.shape
        out_h = (h - 1) // s1 + 1
        out_w = (w - 1) // s1 + 1
        out_k = 2 * d // s2 + 1
        result = torch.zeros(n, out_k ** 2, out_h, out_w, device=x1.device)
        ...

Is this expected to work?

[jamesr66a]

Oh weird, I think this probably has to do with how __torch_function__ works. @ezyang @hameerabbasi does __torch_function__ work on constructor functions like this?

[hameerabbasi]

__torch_function__ only works on functions that take in Tensors. Construction functions don’t work, unless they’re of the kind _like, which is a possible workaround.

[jamesr66a]

@vkuzo does zeros_like work here?

[vkuzo]

this is the real code snippet where we hit this:

    def forward(self, x1, x2):
        assert x1.shape == x2.shape
        d = self.d
        s1 = self.s1
        s2 = self.s2
        n, c, h, w = x1.shape
        out_h = (h - 1) // s1 + 1
        out_w = (w - 1) // s1 + 1
        out_k = 2 * d // s2 + 1
        result = torch.zeros(n, out_k ** 2, out_h, out_w, device=x1.device)
        ...

subbing zeros_like for zeros would need a tensor with the correct shape already constructed. Is there a recommended way to do the tensor construction?

[ezyang]

Oh, just use new_zeros instead of zeros_like, which takes a shape too.

[vkuzo]

Oh, just use new_zeros instead of zeros_like, which takes a shape too.

hmm, unfortunately that doesn't work either:

import torch

class M(torch.nn.Module):
    def forward(self, x):
        x = torch.tensor(()).new_zeros((x.shape[0], x.shape[1]))
        return x
    
m = M()
gm = torch.fx.symbolic_trace(m)
print(gm)

...

<ipython-input-30-78aa572f8f3a> in forward(self, x)
      3 class M(torch.nn.Module):
      4     def forward(self, x):
----> 5         x = torch.tensor(()).new_zeros((x.shape[0], x.shape[1]))
      6         return x
      7 

TypeError: new_zeros(): argument 'size' must be tuple of ints, but found element of type Proxy at pos 1

[hameerabbasi]

I think you should be using x.new_zeros(...) where you forward new_zeros appropriately.

[vkuzo]

I see, thanks, using x.new_zeros(...) instead of torch.tensor(()).new_zeros(...) fixes the symbolic tracing error. Just for my own curiosity, would you have context on why this fixes things?

[hameerabbasi]

Tensor.new_zeros is a classmethod, and is forwarded to the appropriate class, which means it can be traced. Using torch.tensor(()).new_zeros(...) calls the new_zeros classmethod on the Tensor class still, instead of type(x), which then doesn't allow it to be traced anymore (as Tensor itself has no such functionality, and calls to an instance of Tensor or Tensor itself aren't dispatched).

[vkuzo]

Tensor.new_zeros is a classmethod, and is forwarded to the appropriate class, which means it can be traced. Using torch.tensor(()).new_zeros(...) calls the new_zeros classmethod on the Tensor class still, instead of type(x), which then doesn't allow it to be traced anymore (as Tensor itself has no such functionality, and calls to an instance of Tensor or Tensor itself aren't dispatched).

got it, thanks for explaining!