Move test.

Author: ysiraichi

test/test_operations.py

@@ -2367,17 +2367,6 @@ def test_isneginf_no_fallback(self):
     t = t.to(torch.float16)
     self._test_no_fallback(torch.isneginf, (t,))
 
-  def test_trace_raises_error_on_non_matrix_input(self):
-    device = torch_xla.device()
-    a = torch.rand(2, 2, 2, device=device)
-
-    try:
-      torch.trace(a)
-    except RuntimeError as e:
-      expected_error = ("trace(): expected the input tensor f32[2,2,2] to be a "
-                        "matrix (i.e. a 2D tensor).")
-      self.assertEqual(str(e), expected_error)
-
 
 class MNISTComparator(nn.Module):
 

test/test_ops_error_message.py

@@ -222,6 +222,19 @@ def test():
         expect="""roll(): expected `dims` [0] (size=1) to match the size of `shifts` [2, 2] (size=2)."""
     )
 
+  def test_trace_raises_error_on_non_matrix_input(self):
+    device = torch_xla.device()
+    a = torch.rand(2, 2, 2, device=device)
+
+    def test():
+      torch.trace(a)
+
+    self.assertExpectedRaisesInline(
+        exc_type=RuntimeError,
+        callable=test,
+        expect="""trace(): expected the input tensor f32[2,2,2] to be a matrix (i.e. a 2D tensor)."""
+    )
+
 
 if __name__ == "__main__":
   unittest.main()