Reland upstream commit 52cf1ae

[whitneywhtsang]

Reland 52cf1ae that is temporarily reverted in #3021 by c5f5ac1.

...
FAILED language/test_core.py::test_dot[1-64-64-64-4-False-False-softmax-ieee-float32-float32-1] - AssertionError: 
FAILED language/test_core.py::test_dot[1-64-64-64-4-False-False-chain-dot-tf32x3-float32-float32-1] - AssertionError: 
FAILED language/test_core.py::test_dot[1-64-64-64-4-False-False-chain-dot-ieee-float32-float32-1] - AssertionError: 
FAILED language/test_core.py::test_dot[1-32-32-32-4-False-False-none-tf32x3-float32-float32-1] - AssertionError: 
FAILED language/test_core.py::test_dot[1-64-64-64-4-True-True-none-ieee-float32-float32-1] - AssertionError: 
FAILED language/test_core.py::test_dot[1-64-64-64-4-True-False-none-ieee-float32-float32-1] - AssertionError: 
FAILED language/test_core.py::test_dot[1-64-64-64-4-False-True-none-ieee-float32-float32-1] - AssertionError: 
FAILED language/test_core.py::test_dot[1-64-64-64-4-False-False-none-ieee-float32-float32-1_1] - AssertionError: 
======================================================= 46 failed, 11875 passed, 112 skipped, 448 xfailed, 119 warnings in 1370.41s (0:22:50)