Simplify the logic of adjusting signs for estimated eigenvectors comp…

…arisons

Summary: Leverage PyTorch indices comparions to adjust the signs of `estimated_eigenvectors` more Pythonic.

Reviewed By: gajjanag

Differential Revision: D65866712

fbshipit-source-id: 2ec3e50faaf75b4a98027a7dac867c1528f44d72

tests/matrix_functions_test.py

@@ -871,13 +871,10 @@ def test_matrix_eigenvectors(self) -> None:
                         eigenvector_computation_config=qr_config,
                     )
                     # Ensure that the signs of the eigenvectors are consistent.
-                    for col in range(A.shape[1]):
-                        if (
-                            expected_eigenvectors[0, col]
-                            / estimated_eigenvectors[0, col]
-                            < 0
-                        ):
-                            estimated_eigenvectors[:, col] *= -1
+                    estimated_eigenvectors[
+                        :,
+                        expected_eigenvectors[0, :] / estimated_eigenvectors[0, :] < 0,
+                    ] *= -1
                     torch.testing.assert_close(
                         expected_eigenvectors,
                         estimated_eigenvectors,