[Functionalization] Dynamic shape tests

[alanwaketan]

This is the issue that keeps track of the progress being made to fix dynamic shape tests for Functionalization Phase 2.

Here is the list of the failure tests:

• test_simple_expand
• test_simple_expand_on_2d_tensor
• test_forward_pass_dynamic_input_correctness
• test_forward_pass_dynamic_input_compile_once
• test_nonzero_cast
• test_squeeze_nonzero
• test_sizeAdd

CC @wonjoolee95, @vanbasten23

[alanwaketan]

Tried to enable python dispatcher according to pytorch/pytorch#88787 (comment):

diff --git a/test/test_dynamic_shapes.py b/test/test_dynamic_shapes.py
index de6882ae..ecb15b81 100644
--- a/test/test_dynamic_shapes.py
+++ b/test/test_dynamic_shapes.py
@@ -2,23 +2,26 @@ import torch, torch_xla
 import torch_xla.core.xla_model as xm
 import unittest
 
+from torch._dispatch.python import enable_python_dispatcher
+
 dev = xm.xla_device()
 
 
 class TestDynamicShapes(unittest.TestCase):
 
-  @unittest.skip("fails with functionalization")
+  # @unittest.skip("fails with functionalization")
   def test_simple_expand(self):
-    size1 = 5
-    size2 = 2
-    t1 = torch.zeros([size1, size2], device=dev)
-    t1[3][0] = 1
-    # t2 has size [<=10, 2]
-    t2 = torch.nonzero(t1)
-    t5 = torch.ones(1, device=dev)
-    t5.expand(t2.size(0))
-    t5_cpu = t5.cpu()
-    self.assertEqual(t5_cpu.shape[0], 1)
+    with enable_python_dispatcher():
+      size1 = 5
+      size2 = 2
+      t1 = torch.zeros([size1, size2], device=dev)
+      t1[3][0] = 1
+      # t2 has size [<=10, 2]
+      t2 = torch.nonzero(t1)
+      t5 = torch.ones(1, device=dev)
+      t5.expand(t2.size(0))
+      t5_cpu = t5.cpu()
+      self.assertEqual(t5_cpu.shape[0], 1)
 
   @unittest.skip("fails with functionalization")
   def test_simple_expand_on_2d_tensor(self):

Which leads to a different crash:
RuntimeError: NYI

Here is the original crash:
RuntimeError: /workspaces/work/pytorch/build/aten/src/ATen/RegisterCompositeExplicitAutograd.cpp:2109: SymIntArrayRef expected to contain only concrete integers

Investigating the new crash now.

[alanwaketan]

The crash hit:
https://github.com/pytorch/pytorch/blob/master/c10/core/SymNodeImpl.h#L68
The virtual interface for ge.

And then from
https://github.com/pytorch/xla/blob/master/torch_xla/csrc/tensor.cpp#L686
It looks like we indeed haven't implemented ge.

Let me try to implement ge to see if that moves us further. @vanbasten23 @miladm

[alanwaketan]

Fixed ge but then now clone is needed. Not sure how many ops are there. Let's try enable the python dispatcher without functionalization to see if all these missing ops are a cause of functionalization.

[alanwaketan]

Okay, after enabling the python dispatcher on the test without functionalization, the test passes. Therefore, all those missing ops are a cause of functionalization. I guess the reasonable next step is to implement those missing dynamic ops.

[alanwaketan]

After implementing clone, the test passes!!!

[alanwaketan]

#4488 is another attempting fix.

[alanwaketan]

Close as all depending tasks are completed.