mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-06 12:20:52 +01:00
20eac093a7
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/56923 Next Steps in order: - Add backward support for CUDA - Add support for more aggregation types - Benchmarking (for cuda mainly)/more testing/documentation - Support for multi dimension Test Plan: Updated unit test to include 0 length segment as well. Reviewed By: ngimel Differential Revision: D27992228 fbshipit-source-id: 28851811f8a784a63162721c511d69e617a93727
87 lines
2.8 KiB
Python
87 lines
2.8 KiB
Python
import torch
|
|
from torch.testing._internal.common_device_type import (
|
|
instantiate_device_type_tests,
|
|
dtypes,
|
|
dtypesIfCUDA,
|
|
)
|
|
from torch.testing._internal.common_utils import (
|
|
TestCase,
|
|
run_tests,
|
|
gradcheck,
|
|
)
|
|
|
|
|
|
class TestSegmentReductions(TestCase):
|
|
def _test_max_simple_1d(self, device, dtype, unsafe, axis):
|
|
lengths = torch.tensor([1, 2, 3, 0], device=device)
|
|
data = torch.tensor(
|
|
[1, float("nan"), 3, 4, 5, 5],
|
|
device=device,
|
|
dtype=dtype,
|
|
requires_grad=True,
|
|
)
|
|
initial_value = 0
|
|
expected_result = torch.tensor(
|
|
[1, float("nan"), 5, initial_value], device=device, dtype=dtype
|
|
)
|
|
actual_result = torch.segment_reduce(
|
|
data=data,
|
|
reduce="max",
|
|
lengths=lengths,
|
|
axis=axis,
|
|
unsafe=unsafe,
|
|
initial=initial_value,
|
|
)
|
|
self.assertEqual(
|
|
expected_result, actual_result, rtol=1e-03, atol=1e-05, equal_nan=True
|
|
)
|
|
|
|
# Backward is only supported for cpu tensors for now. Return early if cuda
|
|
if data.is_cuda:
|
|
return
|
|
|
|
# Test backward
|
|
expected_grad = torch.tensor([1, 1, 0, 0, 0.5, 0.5], device=device, dtype=dtype)
|
|
actual_result.sum().backward()
|
|
self.assertEqual(
|
|
expected_grad, data.grad, rtol=1e-03, atol=1e-05, equal_nan=True
|
|
)
|
|
|
|
# gradcheck does not work well with bfloat16 or fp16 cpu types
|
|
# also there is small numerical difference with fp32
|
|
if dtype not in [torch.half, torch.bfloat16, torch.float]:
|
|
# gradcheck does not like "nan" input
|
|
data = torch.tensor(
|
|
[1, 10, 3, 4, 5, 5],
|
|
device=device,
|
|
dtype=dtype,
|
|
requires_grad=True,
|
|
)
|
|
self.assertTrue(
|
|
gradcheck(
|
|
lambda x: torch.segment_reduce(
|
|
data=x,
|
|
reduce="max",
|
|
lengths=lengths,
|
|
axis=axis,
|
|
unsafe=unsafe,
|
|
initial=initial_value,
|
|
),
|
|
(data,),
|
|
)
|
|
)
|
|
|
|
@dtypesIfCUDA(torch.half, torch.bfloat16, torch.float, torch.double)
|
|
@dtypes(torch.half, torch.bfloat16, torch.float, torch.double)
|
|
def test_max_simple_1d(self, device, dtype):
|
|
self._test_max_simple_1d(device, dtype, False, 0)
|
|
self._test_max_simple_1d(device, dtype, False, -1)
|
|
self._test_max_simple_1d(device, dtype, True, 0)
|
|
self._test_max_simple_1d(device, dtype, True, -1)
|
|
|
|
|
|
instantiate_device_type_tests(TestSegmentReductions, globals())
|
|
|
|
if __name__ == "__main__":
|
|
run_tests()
|