Repository: pytorch
Issue: Tests stopped working after merging #126905 on Arm64
Labels: module: ci,module: tests,triaged,module: arm,oncall: pt2

### 🐛 Describe the bug

After merging the #126905 PR some tests stopped working on Arm64 (Ampere Altra machine).

One of the tests is:
`python dynamo/test_dynamic_shapes.py -k DynamicShapesFunctionTests.test_pow_int_dynamic_shapes`

and another one is similar from `dynamo/test_functions.py`:
python dynamo/test_functions.py -k FunctionTests.test_pow_int

Here is the full error report
```
python dynamo/test_dynamic_shapes.py -k DynamicShapesFunctionTests.test_pow_int_dynamic_shapes
E0730 11:53:00.945000 281473297682464 torch/fx/experimental/recording.py:280] [0/0] failed while running evaluate_expr(*(False, True), **{'fx_node': eq_4})
F
======================================================================
FAIL: test_pow_int_dynamic_shapes (__main__.DynamicShapesFunctionTests.test_pow_int_dynamic_shapes)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/usr/local/lib/python3.11/site-packages/torch/_dynamo/convert_frame.py", line 564, in transform
    tracer.run()
  File "/usr/local/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 2450, in run
    super().run()
  File "/usr/local/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 892, in run
    while self.step():
          ^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 804, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "/usr/local/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 498, in wrapper
    return inner_fn(self, inst)
           ^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 2058, in CALL
    self.call_function(fn, args, kwargs)
  File "/usr/local/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 742, in call_function
    self.push(fn.call_function(self, args, kwargs))
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/torch/_dynamo/variables/torch.py", line 756, in call_function
    tensor_variable = wrap_fx_proxy(
                      ^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/torch/_dynamo/variables/builder.py", line 1713, in wrap_fx_proxy
    return wrap_fx_proxy_cls(target_cls=TensorVariable, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/torch/_dynamo/variables/builder.py", line 1798, in wrap_fx_proxy_cls
    example_value = get_fake_value(proxy.node, tx, allow_non_graph_fake=True)
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/torch/_dynamo/utils.py", line 1854, in get_fake_value
    raise TorchRuntimeError(str(e)).with_traceback(e.__traceback__) from None
  File "/usr/local/lib/python3.11/site-packages/torch/_dynamo/utils.py", line 1786, in get_fake_value
    ret_val = wrap_fake_exception(
              ^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/torch/_dynamo/utils.py", line 1301, in wrap_fake_exception
    return fn()
           ^^^^
  File "/usr/local/lib/python3.11/site-packages/torch/_dynamo/utils.py", line 1787, in <lambda>
    lambda: run_node(tx.output, node, args, kwargs, nnmodule)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/torch/_dynamo/utils.py", line 1922, in run_node
    raise RuntimeError(make_error_message(e)).with_traceback(
  File "/usr/local/lib/python3.11/site-packages/torch/_dynamo/utils.py", line 1904, in run_node
    return node.target(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/torch/utils/_stats.py", line 20, in wrapper
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/torch/_subclasses/fake_tensor.py", line 1060, in __torch_dispatch__
    return self.dispatch(func, types, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/torch/_subclasses/fake_tensor.py", line 1449, in dispatch
    return self._cached_dispatch_impl(func, types, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/torch/_subclasses/fake_tensor.py", line 1152, in _cached_dispatch_impl
    output = self._dispatch_impl(func, types, args, kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/torch/_subclasses/fake_tensor.py", line 1689, in _dispatch_impl
    return decomposition_table[func](*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/torch/_prims_common/wrappers.py", line 265, in _fn
    result = fn(*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/torch/_prims_common/wrappers.py", line 137, in _fn
    result = fn(**bound.arguments)
             ^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/torch/_refs/__init__.py", line 1043, in _ref
    output = prim(a, b)
             ^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/torch/_refs/__init__.py", line 1200, in pow
    elif b == 2.0:
         ^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/torch/__init__.py", line 562, in __bool__
    return self.node.bool_()
           ^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/torch/fx/experimental/sym_node.py", line 475, in bool_
    return self.guard_bool("", 0)
           ^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/torch/fx/experimental/sym_node.py", line 413, in guard_bool
    r = self.shape_env.evaluate_expr(self.expr, self.hint, fx_node=self.fx_node)
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/torch/fx/experimental/recording.py", line 274, in wrapper
    return event.run(self)
           ^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/torch/fx/experimental/recording.py", line 161, in run
    return self.f(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/torch/fx/experimental/symbolic_shapes.py", line 5131, in evaluate_expr
    assert static_expr == hint, f"{static_expr} != {hint}"
           ^^^^^^^^^^^^^^^^^^^
torch._dynamo.exc.TorchRuntimeError: Failed running call_function <built-in method pow of type object at 0xffff9a9ffd98>(*(FakeTensor(..., size=(s1, s1)), s0), **{}):
False != True

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/usr/local/lib/python3.11/site-packages/torch/testing/_internal/common_utils.py", line 2715, in wrapper
    method(*args, **kwargs)
  File "/usr/local/lib/python3.11/site-packages/torch/_dynamo/testing.py", line 309, in _fn
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/aio/pytorch-dls/test/dynamo/test_functions.py", line 1789, in test_pow_int
    self.assertEqual(opt_fn(x, 2), fn(x, 2))
                     ^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/torch/_dynamo/eval_frame.py", line 421, in _fn
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/torch/_dynamo/convert_frame.py", line 1078, in catch_errors
    return callback(frame, cache_entry, hooks, frame_state, skip=1)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/torch/_dynamo/convert_frame.py", line 456, in _convert_frame_assert
    return _compile(
           ^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/torch/_utils_internal.py", line 83, in wrapper_function
    return StrobelightCompileTimeProfiler.profile_compile_time(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/torch/_strobelight/compile_time_profiler.py", line 129, in profile_compile_time
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/contextlib.py", line 81, in inner
    return func(*args, **kwds)
           ^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/torch/_dynamo/convert_frame.py", line 799, in _compile
    guarded_code = compile_inner(code, one_graph, hooks, transform)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/torch/_dynamo/utils.py", line 232, in time_wrapper
    r = func(*args, **kwargs)
        ^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/torch/_dynamo/convert_frame.py", line 618, in compile_inner
    out_code = transform_code_object(code, transform)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/torch/_dynamo/bytecode_transformation.py", line 1184, in transform_code_object
    transformations(instructions, code_options)
  File "/usr/local/lib/python3.11/site-packages/torch/_dynamo/convert_frame.py", line 177, in _fn
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/torch/_dynamo/convert_frame.py", line 572, in transform
    bisect(tracer.output.shape_env)
  File "/usr/local/lib/python3.11/site-packages/torch/fx/experimental/validator.py", line 756, in bisect
    assert left in exception and isinstance(exception[left], ValidationException)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
AssertionError: 


You can suppress this exception and fall back to eager by setting:
    import torch._dynamo
    torch._dynamo.config.suppress_errors = True


To execute this test, run the following from the base repo dir:
     python test/dynamo/test_dynamic_shapes.py -k DynamicShapesFunctionTests.test_pow_int_dynamic_shapes

This message can be suppressed by setting PYTORCH_PRINT_REPRO_ON_FAILURE=0

----------------------------------------------------------------------
Ran 1 test in 0.398s

FAILED (failures=1)
```

Overall I see there is a lot of tests failing after building clean Pytorch 2.4 build, which haven't existed on 2.2.1 on Ampere Altra machine. It feels like CI isn't testing ARM64 builds well. I tried to find an ARM build on Github actions and this found job: https://github.com/pytorch/pytorch/actions/workflows/linux-aarch64.yml
However it runs only few selected tests:
`python test/run_test.py --include test_modules test_mkldnn test_mkldnn_fusion test_openmp test_torch test_dynamic_shapes test_transformers test_multiprocessing test_numpy_interop --verbose`, not the whole test suite
 
@ezyang 


### Versions

Collecting environment information...
PyTorch version: 2.4.0a0+git935a946
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.3 LTS (aarch64)
GCC version: (Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0
Clang version: 14.0.0-1ubuntu1.1
CMake version: version 3.22.1
Libc version: glibc-2.35

Python version: 3.11.9 (main, Jun 13 2024, 04:17:51) [GCC 12.2.0] (64-bit runtime)
Python platform: Linux-5.15.0-100-generic-aarch64-with-glibc2.35
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:                       aarch64
CPU op-mode(s):                     32-bit, 64-bit
Byte Order:                         Little Endian
CPU(s):                             160
On-line CPU(s) list:                0-159
Vendor ID:                          ARM
Model name:                         Neoverse-N1
Model:                              1
Thread(s) per core:                 1
Core(s) per socket:                 80
Socket(s):                          2
Stepping:                           r3p1
Frequency boost:                    disabled
CPU max MHz:                        3000.0000
CPU min MHz:                        1000.0000
BogoMIPS:                           50.00
Flags:                              fp asimd evtstrm aes pmull sha1 sha2 crc32 atomics fphp asimdhp cpuid asimdrdm lrcpc dcpop asimddp ssbs
L1d cache:                          10 MiB (160 instances)
L1i cache:                          10 MiB (160 instances)
L2 cache:                           160 MiB (160 instances)
NUMA node(s):                       2
NUMA node0 CPU(s):                  0-79
NUMA node1 CPU(s):                  80-159
Vulnerability Gather data sampling: Not affected
Vulnerability Itlb multihit:        Not affected
Vulnerability L1tf:                 Not affected
Vulnerability Mds:                  Not affected
Vulnerability Meltdown:             Not affected
Vulnerability Mmio stale data:      Not affected
Vulnerability Retbleed:             Not affected
Vulnerability Spec rstack overflow: Not affected
Vulnerability Spec store bypass:    Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1:           Mitigation; __user pointer sanitization
Vulnerability Spectre v2:           Mitigation; CSV2, BHB
Vulnerability Srbds:                Not affected
Vulnerability Tsx async abort:      Not affected

Versions of relevant libraries:
[pip3] mypy==1.4.1
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.24.4
[pip3] onnx==1.16.1
[pip3] onnxruntime==1.18.0
[pip3] optree==0.11.0
[pip3] torch==2.4.0a0+gitfb696ef
[pip3] torchvision==0.19.0a0+48b1edf
[conda] Could not collect

cc @seemethere @malfet @pytorch/pytorch-dev-infra @mruberry @ZainRizvi @snadampal @milpuz01 @ezyang @chauhang @penguinwu