[MPS] Add support for nansum on mps #93845
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/93845
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit e921ed1: This comment was automatically generated by Dr. CI and updates every 15 minutes. |
pytorch-bot
bot
added
ciflow/mps
test/test_mps.py
Outdated
| def helper(self, dtype, noncontiguous, dim): | ||
| zero_cpu = torch.zeros((), dtype=dtype) | ||
|
|
||
| # Randomly scale the values | ||
| scale = random.randint(10, 100) | ||
| x_cpu: torch.Tensor = make_tensor( | ||
| (5, 5), dtype=dtype, device='cpu', | ||
| low=-scale, high=scale, noncontiguous=noncontiguous) | ||
|
|
||
| if dtype.is_floating_point: | ||
| nan_mask_cpu = x_cpu < (0.2 * scale) | ||
| x_no_nan_cpu = torch.where(nan_mask_cpu, zero_cpu, x_cpu) | ||
| x_cpu[nan_mask_cpu] = np.nan | ||
| else: | ||
| x_no_nan_cpu = x_cpu | ||
|
|
||
| x_mps = x_cpu.to('mps') | ||
| actual_out_mps = torch.empty(0, dtype=dtype, device='mps') | ||
| expect_out_cpu = torch.empty(0, dtype=dtype) | ||
| dim_kwargs = {"dim": dim} if dim is not None else {} | ||
| expect = torch.sum(x_no_nan_cpu, **dim_kwargs) | ||
|
|
||
| actual_cpu = torch.nansum(x_cpu, **dim_kwargs) | ||
| # Sanity check on CPU | ||
| self.assertEqual(expect, actual_cpu) | ||
|
|
||
| # Test MPS | ||
| actual_mps = torch.nansum(x_mps, **dim_kwargs) | ||
| # Test out= variant | ||
| torch.nansum(x_mps, out=actual_out_mps, **dim_kwargs) | ||
| torch.nansum(x_cpu, out=expect_out_cpu, **dim_kwargs) | ||
| self.assertEqual(expect, actual_mps) | ||
| self.assertEqual(expect_out_cpu, actual_out_mps) | ||
|
|
||
| def test_nansum(self): | ||
| args = itertools.product( | ||
| (torch.float16, torch.float32, torch.int32, torch.int64), # dtype | ||
| (True, False), # noncontiguous | ||
| (0, 1, None), # dim | ||
| ) | ||
|
|
||
| for dtype, noncontiguous, dim in args: | ||
| with self.subTest(dtype=dtype, noncontiguous=noncontiguous, dim=dim): | ||
| self.helper(dtype, noncontiguous, dim) | ||
|
|
There was a problem hiding this comment.
Please don't make it a new class, but rather add as a method to say test_consistency (and actually, can't it just be added using opinfo?)
There was a problem hiding this comment.
Not exactly sure how to add it with opinfo, I might be missing something, could you give me some guidance?
I tried running EXPECTTEST_ACCEPT=1 python test_mps.py TestConsistencyCPU per the instructions in TestConsistency and the interpreter crashes on unrelated things (e.g. unssuported type conversions like cdouble)
As a default I'll update the PR and move test_nansum into TestMps as a method, but helping me test this with opinfo would be very appreciated :)
* Add nansum_out_mps and nansum_mps functions * Moved `get_dtype_from_self` into ReduceOpsUtils.h
soof-golan
force-pushed
the
feature/nansum-mps
branch
from
8eb45c0 to
e921ed1
Compare
|
@pytorchmergebot merge |
Merge startedYour change will be merged once all checks pass (ETA 0-4 Hours). Learn more about merging in the wiki. Questions? Feedback? Please reach out to the PyTorch DevX Team |
4 participants
nansum_out_mpsandnansum_mpsfunctionsget_dtype_from_selfinto ReduceOpsUtils.hFixes #86809