mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-06 12:20:52 +01:00
e27740b38e
Summary: This is to setup boiler plate code for backward and CPU implementation. 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 also check correctness of backward. Wait for CI signal Reviewed By: ngimel Differential Revision: D27970340 fbshipit-source-id: 3e608c7fe3628b0a761dd8affc6aad8f65a6ef7f
74 lines
2.5 KiB
Python
74 lines
2.5 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], device=device)
|
|
data = torch.tensor(
|
|
[1, float("nan"), 3, 4, 5, 5],
|
|
device=device,
|
|
dtype=dtype,
|
|
requires_grad=True,
|
|
)
|
|
expected_result = torch.tensor([1, float("nan"), 5], device=device, dtype=dtype)
|
|
actual_result = torch.segment_reduce(
|
|
data=data, reduce="max", lengths=lengths, axis=axis, unsafe=unsafe
|
|
)
|
|
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
|
|
),
|
|
(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()
|