TPU+GPU test flake: test_diagonal_write_transposed_r3 (__main__.TestAtenXlaTensor)

[tengyifei]

Seen in https://github.com/pytorch/xla/actions/runs/14486615445/job/40635200589

======================================================================
FAIL: test_diagonal_write_transposed_r3 (__main__.TestAtenXlaTensor)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/home/runner/_work/xla/xla/pytorch/xla/test/test_operations.py", line 1399, in test_diagonal_write_transposed_r3
    self.runAtenTest([torch.randn(5, 8, 7)], test_fn)
  File "/home/runner/_work/xla/xla/pytorch/xla/test/test_utils.py", line 393, in runAtenTest
    self.compareResults(results, xla_results, rel_err=rel_err, abs_err=abs_err)
  File "/home/runner/_work/xla/xla/pytorch/xla/test/test_utils.py", line 378, in compareResults
    self.assertEqualRel(
  File "/home/runner/_work/xla/xla/pytorch/xla/test/test_utils.py", line 332, in assertEqualRel
    self.fail('Relative error higher than the maximum tolerance')
AssertionError: Relative error higher than the maximum tolerance

----------------------------------------------------------------------
Ran 211 tests in 62.095s

FAILED (failures=1, skipped=27)
a=1.64124596118927	b=2.545245885848999	diff=0.903999924659729	index=(1, 1, 0)
a=1.0509189367294312	b=1.9549188613891602	diff=0.903999924659729	index=(2, 1, 1)
a=0.24387085437774658	b=1.1478707790374756	diff=0.903999924659729	index=(3, 1, 2)
a=-0.4914337992668152	b=0.4125661849975586	diff=0.9039999842643738	index=(4, 1, 3)

max_diff=0.9039999842643738	max_rel=1.8395153284072876	index=(4, 1, 3)

[tengyifei]

Seen again: https://github.com/pytorch/xla/actions/runs/14766475918/job/41460021031?pr=9045

[tengyifei]

Fails on GPU too: https://github.com/pytorch/xla/actions/runs/14913309223/job/41897244327?pr=9065