已合并
test: add isclose equal_nan broadcast regression test #39508
21xiu创建于 6月29日
test: add isclose equal_nan broadcast regression test #39508
已合并
共 1 个文件变更+31-0
| @@ -278,6 +278,37 @@ class TestTesting(TestCase): | |||
| 278 | 278 | ||
| 279 | self._isclose_helper(tests, device, dtype, equal_nan=True, rtol=0, atol=0) | 279 | self._isclose_helper(tests, device, dtype, equal_nan=True, rtol=0, atol=0) |
| 280 | 280 | ||
| 281 | + | ||
| 282 | + | ||
| 283 | + def test_isclose_equal_nan_broadcast(self, device, dtype): | ||
| 284 | + # Regression test: isclose with equal_nan=True should handle broadcasting | ||
| 285 | + nan = torch.nan | ||
| 286 | + | ||
| 287 | + # One-sided broadcast: different sizes | ||
| 288 | + a = torch.tensor([nan], device=device, dtype=dtype) | ||
| 289 | + b = torch.tensor([nan, nan], device=device, dtype=dtype) | ||
| 290 | + result = torch.isclose(a, b, equal_nan=True) | ||
| 291 | + self.assertEqual(result, torch.tensor([True, True], device=device)) | ||
| 292 | + | ||
| 293 | + # One-sided broadcast: scalar-like vs vector | ||
| 294 | + a = torch.tensor([nan], device=device, dtype=dtype) | ||
| 295 | + b = torch.tensor([nan, 1.0, nan], device=device, dtype=dtype) | ||
| 296 | + result = torch.isclose(a, b, equal_nan=True) | ||
| 297 | + self.assertEqual(result, torch.tensor([True, False, True], device=device)) | ||
| 298 | + | ||
| 299 | + # Mutual broadcast | ||
| 300 | + a = torch.tensor([[nan], [1.0]], device=device, dtype=dtype) # [2, 1] | ||
| 301 | + b = torch.tensor([[nan, 1.0]], device=device, dtype=dtype) # [1, 2] | ||
| 302 | + result = torch.isclose(a, b, equal_nan=True) | ||
| 303 | + expected = torch.tensor([[True, False], [False, True]], device=device) | ||
| 304 | + self.assertEqual(result, expected) | ||
| 305 | + | ||
| 306 | + # Same shape (fast path) still works | ||
| 307 | + a = torch.tensor([nan, 1.0], device=device, dtype=dtype) | ||
| 308 | + b = torch.tensor([nan, 1.0], device=device, dtype=dtype) | ||
| 309 | + result = torch.isclose(a, b, equal_nan=True) | ||
| 310 | + self.assertEqual(result, torch.tensor([True, True], device=device)) | ||
| 311 | + | ||
| 281 | 312 | ||
| 282 | 313 | ||
| 283 | 314 | ||