Openxla Dynamo backend throw a warning message about boxed argument

[JackCaoG]

🐛 Bug

currently with the new openxla backend it will throw a

/usr/local/lib/python3.10/dist-packages/torch/_functorch/aot_autograd.py:1511: UserWarning: Your compiler for AOTAutograd is returning a function that doesn't take boxed arguments. Please wrap it with functorch.compile.make_boxed_func or handle the boxed arguments yourself. See https://github.com/pytorch/pytorch/pull/83137#issuecomment-1211320670 for rationale.

warning.

[JackCaoG]

I tried to fix it with

--- a/torch/_dynamo/backends/torchxla.py
+++ b/torch/_dynamo/backends/torchxla.py
@@ -4,6 +4,8 @@ import warnings
 from ..backends.common import aot_autograd
 from ..backends.registry import register_experimental_backend as register_backend
 
+from functorch.compile import make_boxed_func
+
 log = logging.getLogger(__name__)
 
 
@@ -24,7 +26,6 @@ def torchxla_trace_once(model, fake_tensor_inputs):
 def openxla_eval(model, fake_tensor_inputs):
     return xla_backend_helper(model, fake_tensor_inputs)
 
-
 def xla_backend_helper(model, fake_tensor_inputs):
     import torch_xla.core.dynamo_bridge as bridge  # type: ignore[import]
 
@@ -38,7 +39,7 @@ def xla_backend_helper(model, fake_tensor_inputs):
             del model
         return compiled_graph(*args)
 
-    return fwd
+    return make_boxed_func(fwd)

need to run some model to check if this works. Error message is gone through.

[JackCaoG]

with above change, using the old torchxla_trace_once I saw the error

  File "/src/llama/llama/generation.py", line 303, in generate
    = self._generate_one_token_fn(
  File "/src/pytorch/torch/_dynamo/eval_frame.py", line 321, in _fn
    return fn(*args, **kwargs)
  File "/src/llama/llama/generation.py", line 197, in _generate_one_token
    def _generate_one_token(self, tokens, input_tokens, input_text_mask,
  File "/src/pytorch/torch/_dynamo/eval_frame.py", line 321, in _fn
    return fn(*args, **kwargs)
  File "/src/pytorch/torch/_dynamo/external_utils.py", line 17, in inner
    return fn(*args, **kwargs)
TypeError: g() takes 1 positional argument but 73 were given

seems like it only expect one boxed argument

[JackCaoG]

switching to openxla backend solved this issue. I guess the take away is that make_boxed_func should be applied to aotbackend only. Non-aot backend should not use it. Let me try to fix it..

[JackCaoG]

This is fixed, closing the issue.