[Dynamic Shape] Reshape does not work with dynamic input + symint target size

[JackCaoG]

🐛 Bug

I would expect below code to emit some meaningful error message since we haven't support reshape.symint(in fact I don't know if such op existed).

import torch
import torch_xla
import torch_xla.core.xla_model as xm

import os
os.environ["XLA_EXPERIMENTAL"] = "nonzero:masked_select:masked_scatter"

dev = xm.xla_device()
size1 = 5
size2 = 2
t1 = torch.zeros([size1, size2], device=dev)
t1[3][0] = 1
# t2 has size [<=10, 2]
t2 = torch.nonzero(t1)
t3 = t2.reshape([t2.size(0) * t2.size(1)])

However I get

Traceback (most recent call last):
  File "test/test_dynamic_fail.py", line 16, in <module>
    t3 = t2.reshape([t2.size(0) * t2.size(1)])
RuntimeError: RuntimeError: NYI

which I am not sure what's the real error under the hood.

@miladm @vanbasten23 @Krovatkin @wconstab

[JackCaoG]

I chatted with Brian, Reshape is also a view op so we can delay the invesgation of this issue after functionalization pass migration is completed.

[vanbasten23]

Thanks for the headsup. Btw, what do you mean by a view op?

[JackCaoG]

https://pytorch.org/docs/stable/tensor_view.html

[miladm]

• Looks like we are blocked on functionalization support before this PR can land. Correct?
• @vanbasten23 can you please investigate if the backward pass requires reshape ops; if so, this issue becomes high priority.

Marking this issue as blocked for the time being.

CC @bdhirsh

[bdhirsh]

@JackCaoG just confirming:

(1) reshape() is a "maybe view" op: it's a composite op, that decomposes into either view or view + clone depending on the striding of the input. So functionalization never operates on reshape() directly - it operates on view()

(2) In my head, "XLA using functionalization" and "XLA getting dynamic shapes working" should be pretty independent. Is the dependency here mostly just organizational? (view ops are in flux so easier to wait till functionalization lands). Or is it somethign else.

[JackCaoG]

In my head, "XLA using functionalization" and "XLA getting dynamic shapes working" should be pretty independent

I agree. My thinking is we want to make view ops to support dynamic shape once functionalization pass is over since it would be easier that way.

[alanwaketan]

@vanbasten23 LMK if I wrongly assigned the bug to you.