[BE]: Apply PLR1736 fixes (unnecessary index lookup) (#127716)
Applies the PLR1736 preview rule with some more autofixes to cut down on unnecessary accesses. Added a noqa since that test actually testing the dunder method.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/127716
Approved by: https://github.com/ezyang
diff --git a/test/distributed/test_device_mesh.py b/test/distributed/test_device_mesh.py
index 03457de..c024ddc 100644
--- a/test/distributed/test_device_mesh.py
+++ b/test/distributed/test_device_mesh.py
@@ -672,9 +672,9 @@
)
unpadded_list = [
(
- unpad_tensor(big_tensor_chunks[i], shard_dim, pad_sizes[i])
+ unpad_tensor(big_tensor, shard_dim, pad_sizes[i])
if pad_sizes[i] > 0
- else big_tensor_chunks[i]
+ else big_tensor
)
for i, big_tensor in enumerate(big_tensor_chunks)
]
diff --git a/test/test_torch.py b/test/test_torch.py
index 717943f..ff57370 100644
--- a/test/test_torch.py
+++ b/test/test_torch.py
@@ -8397,7 +8397,7 @@
def test_iter(self) -> None:
x = torch.randn(5, 5)
for i, sub in enumerate(x):
- self.assertEqual(sub, x[i])
+ self.assertEqual(sub, x[i]) # noqa: PLR1736
x = torch.tensor([])
self.assertEqual(list(x), [])
diff --git a/torch/_decomp/decompositions.py b/torch/_decomp/decompositions.py
index 5bec539..76599d2 100644
--- a/torch/_decomp/decompositions.py
+++ b/torch/_decomp/decompositions.py
@@ -2147,7 +2147,7 @@
def _broadcast_batch_norm_backward(x, broadcast_mask):
for axis, mask in enumerate(broadcast_mask):
- if mask == 1 and not (axis < x.ndim and x.shape[axis] == broadcast_mask[axis]):
+ if mask == 1 and not (axis < x.ndim and x.shape[axis] == mask):
x = x.unsqueeze(axis)
return x
