frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-20 21:14:14 +08:00
Code Issues Packages Projects Releases Wiki Activity

Add Half support for addcmul, addcdiv, cumsum, and topk on CPU (#103319)

Browse Source 
Add Half support for addcmul, addcdiv, cumsum, and topk on CPU.
Note: This PR will introduce the issue  https://github.com/pytorch/pytorch/issues/111454.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/103319
Approved by: https://github.com/jgong5, https://github.com/cpuhrsch
This commit is contained in:
CaoE
2023-10-19 17:47:45 +00:00
committed by PyTorch MergeBot
parent 715dfced72
commit 2a40b7efcb
8 changed files with 63 additions and 60 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
test/test_sort_and_select.py
View File

@ -754,9 +754,8 @@ class TestSortAndSelect(TestCase):
for curr_size in (small, large, verylarge):
self._test_topk_dtype(device, dtype, True, curr_size)
@onlyCUDA
@dtypes(torch.bfloat16)
def test_topk_bfloat16(self, device, dtype):
@dtypes(torch.bfloat16, torch.half)
def test_topk_lower_precision(self, device, dtype):
small = 10
large = 4096
@ -765,7 +764,7 @@ class TestSortAndSelect(TestCase):
self._test_topk_dtype(device, dtype, False, curr_size)
@dtypesIfCUDA(*floating_types_and(torch.half, torch.bfloat16))
@dtypes(torch.float, torch.double, torch.bfloat16)
@dtypes(torch.float, torch.double, torch.bfloat16, torch.half)
def test_topk_nonfinite(self, device, dtype):
x = torch.tensor([float('nan'), float('inf'), 1e4, 0, -1e4, -float('inf')], device=device, dtype=dtype)
val, idx = x.topk(4)
@ -796,7 +795,7 @@ class TestSortAndSelect(TestCase):
@onlyNativeDeviceTypes
@dtypesIfCUDA(*all_types_and(torch.bfloat16))
@dtypes(*all_types())
@dtypes(*all_types_and(torch.bfloat16, torch.half))
def test_topk_zero(self, device, dtype):
# https://github.com/pytorch/pytorch/issues/49205
t = torch.rand(2, 2, device=device).to(dtype=dtype)