[test] Fori loop simple case without hard-code

[ManfeiBai]

code based on #7012, create post-order tracing interface on python level, to enable post-order tracing for XLAComputation like https://github.com/pytorch/xla/blob/master/torch_xla/csrc/xla_graph_executor.cpp#L1293

[ManfeiBai]

Hi, @miladm, this is the script code for post-order tracing compile function

[ManfeiBai]

Hi, @miladm, this is the script code implementation for post-order tracing compile function

[JackCaoG]

I have

  def test_while_loop_tpu_simple_linear(self):

    xm.mark_step()
    device = xm.xla_device()
    #device = ''
    torch.set_grad_enabled(False)

    class SimpleWithLinear(torch.nn.Module):
      def __init__(self):
          super().__init__()
          self.linear = torch.nn.Linear(2, 2)
          self.register_buffer("dec", torch.tensor(1))

      def forward(self, iter, x):
          def cond_fn(it, x):
              return it - self.dec > 0

          def body_fn(it, x):
              return it - 1, self.linear(x)

          return while_loop(cond_fn, body_fn, (iter, x))
      
    simple_with_linear = SimpleWithLinear()
    simple_with_linear.to(device)
    #breakpoint()
    input = torch.randn(2, 2).to(device)
    iter = torch.tensor(3, device=device)
    res = simple_with_linear(iter, input)

which works on CPU device but error out for XLA devices. Let's try to get this simple example working.

[ManfeiBai]

script used for compare order of upper/lower for post-order

[ManfeiBai]

thanks, based on this PR, created and updated to simple test case in #7094, would close this pr for now