wehub-resource-sync
318 lines
11 KiB
Python
318 lines
11 KiB
Python
# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
"""Tests for paddle.nansum PHI kernel implementation."""
|
|
|
|
import unittest
|
|
|
|
import numpy as np
|
|
|
|
import paddle
|
|
|
|
|
|
def np_nansum(x, axis=None, keepdims=False, dtype=None):
|
|
"""Reference implementation using numpy."""
|
|
if dtype is not None:
|
|
return np.nansum(x, axis=axis, keepdims=keepdims).astype(dtype)
|
|
return np.nansum(x, axis=axis, keepdims=keepdims)
|
|
|
|
|
|
def np_nansum_grad(x, out_grad_broadcast):
|
|
"""Reference grad: broadcast(out_grad) masked by ~isnan(x)."""
|
|
grad = out_grad_broadcast.copy()
|
|
grad[np.isnan(x)] = 0.0
|
|
return grad
|
|
|
|
|
|
class TestNansumForward(unittest.TestCase):
|
|
"""Test nansum forward correctness on various cases."""
|
|
|
|
def setUp(self):
|
|
self.places = ['cpu']
|
|
if paddle.is_compiled_with_cuda():
|
|
self.places.append('gpu')
|
|
|
|
def _run_test(self, x_np, axis=None, keepdim=False, dtype=None):
|
|
expected = np_nansum(x_np, axis=axis, keepdims=keepdim, dtype=dtype)
|
|
paddle_dtype = None
|
|
if dtype == 'float64':
|
|
paddle_dtype = paddle.float64
|
|
elif dtype == 'float32':
|
|
paddle_dtype = paddle.float32
|
|
for place in self.places:
|
|
paddle.device.set_device(str(place))
|
|
x = paddle.to_tensor(x_np, place=place)
|
|
out = paddle.nansum(
|
|
x, axis=axis, keepdim=keepdim, dtype=paddle_dtype
|
|
)
|
|
np.testing.assert_allclose(
|
|
out.numpy(),
|
|
expected,
|
|
rtol=1e-5,
|
|
atol=1e-6,
|
|
err_msg=f"Failed on {place}, axis={axis}, keepdim={keepdim}",
|
|
)
|
|
|
|
def test_all_nan(self):
|
|
"""nansum of all-NaN tensor should be 0."""
|
|
x = np.array(
|
|
[float('nan'), float('nan'), float('nan')], dtype='float32'
|
|
)
|
|
self._run_test(x)
|
|
|
|
def test_no_nan(self):
|
|
"""nansum without NaN should equal sum."""
|
|
x = np.array([1.0, 2.0, 3.0, 4.0], dtype='float32')
|
|
self._run_test(x)
|
|
|
|
def test_mixed_nan(self):
|
|
"""Basic mixed NaN/value test."""
|
|
x = np.array(
|
|
[[float('nan'), 0.3, 0.5, 0.9], [0.1, 0.2, float('nan'), 0.7]],
|
|
dtype='float32',
|
|
)
|
|
self._run_test(x)
|
|
self._run_test(x, axis=0)
|
|
self._run_test(x, axis=1)
|
|
self._run_test(x, axis=-1)
|
|
|
|
def test_keepdim(self):
|
|
x = np.array(
|
|
[[float('nan'), 1.0], [2.0, float('nan')]], dtype='float32'
|
|
)
|
|
self._run_test(x, axis=1, keepdim=True)
|
|
self._run_test(x, axis=0, keepdim=True)
|
|
|
|
def test_reduce_all(self):
|
|
"""axis=None reduces all dims."""
|
|
x = np.array(
|
|
[[[1, float('nan')], [3, 4]], [[5, 6], [float('nan'), 8]]],
|
|
dtype='float32',
|
|
)
|
|
self._run_test(x)
|
|
|
|
def test_multi_axis(self):
|
|
x = np.array(
|
|
[[[1, float('nan')], [3, 4]], [[5, 6], [float('nan'), 8]]],
|
|
dtype='float32',
|
|
)
|
|
self._run_test(x, axis=(1, 2))
|
|
self._run_test(x, axis=(0, 1))
|
|
|
|
def test_dtype_promotion(self):
|
|
"""Test output dtype control."""
|
|
x = np.array([1.0, float('nan'), 3.0], dtype='float32')
|
|
self._run_test(x, dtype='float64')
|
|
|
|
def test_integer_input(self):
|
|
"""Integer types have no NaN; nansum == sum."""
|
|
x = np.array([1, 2, 3, 4], dtype='int32')
|
|
self._run_test(x)
|
|
self._run_test(x, axis=0)
|
|
|
|
def test_empty_tensor(self):
|
|
"""nansum of empty tensor should be 0."""
|
|
for place in self.places:
|
|
paddle.device.set_device(str(place))
|
|
x = paddle.empty([0, 3], dtype='float32')
|
|
out = paddle.nansum(x)
|
|
self.assertEqual(out.item(), 0.0)
|
|
|
|
def test_empty_tensor_int64(self):
|
|
"""nansum of empty int32 tensor with dtype=int64 should be 0."""
|
|
for place in self.places:
|
|
paddle.device.set_device(str(place))
|
|
x = paddle.empty([0, 3], dtype='int32')
|
|
out = paddle.nansum(x, dtype=paddle.int64)
|
|
self.assertEqual(out.item(), 0)
|
|
self.assertEqual(out.dtype, paddle.int64)
|
|
|
|
def test_neg_nan(self):
|
|
"""-NaN should also be treated as 0."""
|
|
x = np.array([1.0, float('-nan'), 3.0], dtype='float32')
|
|
self._run_test(x)
|
|
|
|
def test_single_element(self):
|
|
x_nan = np.array([float('nan')], dtype='float32')
|
|
x_val = np.array([5.0], dtype='float32')
|
|
self._run_test(x_nan)
|
|
self._run_test(x_val)
|
|
|
|
|
|
class TestNansumBackward(unittest.TestCase):
|
|
"""Test nansum backward (gradient) correctness."""
|
|
|
|
def setUp(self):
|
|
self.places = ['cpu']
|
|
if paddle.is_compiled_with_cuda():
|
|
self.places.append('gpu')
|
|
|
|
def _check_grad(self, x_np, axis=None, keepdim=False):
|
|
expected_out = np_nansum(x_np, axis=axis, keepdims=keepdim)
|
|
# Compute expected gradient: ones broadcast to x shape, masked by ~isnan
|
|
grad_out = np.ones_like(expected_out)
|
|
# Broadcast grad_out to x shape
|
|
if axis is not None:
|
|
if isinstance(axis, int):
|
|
axes = [axis]
|
|
else:
|
|
axes = list(axis)
|
|
# Normalize negative axes
|
|
axes = [a % x_np.ndim for a in axes]
|
|
expand_shape = list(x_np.shape)
|
|
for a in axes:
|
|
expand_shape[a] = 1
|
|
grad_broadcast = grad_out.reshape(expand_shape)
|
|
grad_broadcast = np.broadcast_to(grad_broadcast, x_np.shape)
|
|
else:
|
|
grad_broadcast = np.broadcast_to(grad_out, x_np.shape)
|
|
expected_grad = np_nansum_grad(x_np, grad_broadcast)
|
|
|
|
for place in self.places:
|
|
paddle.device.set_device(str(place))
|
|
x = paddle.to_tensor(x_np, place=place, stop_gradient=False)
|
|
out = paddle.nansum(x, axis=axis, keepdim=keepdim)
|
|
out.backward()
|
|
np.testing.assert_allclose(
|
|
x.grad.numpy(),
|
|
expected_grad,
|
|
rtol=1e-5,
|
|
atol=1e-6,
|
|
err_msg=f"Grad failed on {place}, axis={axis}",
|
|
)
|
|
|
|
def test_grad_basic(self):
|
|
x = np.array(
|
|
[[float('nan'), 0.3, 0.5, 0.9], [0.1, 0.2, float('nan'), 0.7]],
|
|
dtype='float32',
|
|
)
|
|
self._check_grad(x)
|
|
|
|
def test_grad_axis0(self):
|
|
x = np.array(
|
|
[[float('nan'), 1.0], [2.0, float('nan')]], dtype='float32'
|
|
)
|
|
self._check_grad(x, axis=0)
|
|
|
|
def test_grad_axis1(self):
|
|
x = np.array(
|
|
[[float('nan'), 1.0], [2.0, float('nan')]], dtype='float32'
|
|
)
|
|
self._check_grad(x, axis=1)
|
|
|
|
def test_grad_all_nan(self):
|
|
"""All-NaN: gradient should be all zeros."""
|
|
x = np.array([float('nan'), float('nan')], dtype='float32')
|
|
self._check_grad(x)
|
|
|
|
def test_grad_no_nan(self):
|
|
"""No NaN: gradient should be all ones (like sum)."""
|
|
x = np.array([1.0, 2.0, 3.0], dtype='float32')
|
|
self._check_grad(x)
|
|
|
|
def test_grad_keepdim(self):
|
|
x = np.array([[float('nan'), 1.0], [2.0, 3.0]], dtype='float32')
|
|
self._check_grad(x, axis=1, keepdim=True)
|
|
|
|
def test_grad_3d_multi_axis(self):
|
|
x = np.array(
|
|
[[[1, float('nan')], [3, 4]], [[5, 6], [float('nan'), 8]]],
|
|
dtype='float32',
|
|
)
|
|
self._check_grad(x, axis=(1, 2))
|
|
|
|
def test_grad_float64(self):
|
|
x = np.array([float('nan'), 1.0, 2.0], dtype='float64')
|
|
self._check_grad(x)
|
|
|
|
def test_grad_empty_tensor(self):
|
|
"""Backward on empty tensor: x_grad should be empty with correct shape."""
|
|
for place in self.places:
|
|
paddle.device.set_device(str(place))
|
|
x = paddle.empty([0, 3], dtype='float32')
|
|
x.stop_gradient = False
|
|
out = paddle.nansum(x)
|
|
out.backward()
|
|
self.assertEqual(list(x.grad.shape), [0, 3])
|
|
|
|
|
|
class TestNansumAlignPyTorch(unittest.TestCase):
|
|
"""Explicit alignment tests with known PyTorch outputs."""
|
|
|
|
def setUp(self):
|
|
self.places = ['cpu']
|
|
if paddle.is_compiled_with_cuda():
|
|
self.places.append('gpu')
|
|
|
|
def test_torch_example_1(self):
|
|
"""torch.nansum(tensor([nan, 0.3, 0.5, 0.9, 0.1, 0.2, nan, 0.7])) = 2.7"""
|
|
x_np = np.array(
|
|
[float('nan'), 0.3, 0.5, 0.9, 0.1, 0.2, float('nan'), 0.7],
|
|
dtype='float32',
|
|
)
|
|
for place in self.places:
|
|
paddle.device.set_device(str(place))
|
|
x = paddle.to_tensor(x_np, place=place)
|
|
out = paddle.nansum(x)
|
|
np.testing.assert_allclose(out.numpy(), 2.7, rtol=1e-5)
|
|
|
|
def test_torch_example_2d_axis0(self):
|
|
"""Matches torch.nansum(x, dim=0) for 2x4 with NaN."""
|
|
x_np = np.array(
|
|
[[float('nan'), 0.3, 0.5, 0.9], [0.1, 0.2, float('-nan'), 0.7]],
|
|
dtype='float32',
|
|
)
|
|
expected = np.array([0.1, 0.5, 0.5, 1.6], dtype='float32')
|
|
for place in self.places:
|
|
paddle.device.set_device(str(place))
|
|
x = paddle.to_tensor(x_np, place=place)
|
|
out = paddle.nansum(x, axis=0)
|
|
np.testing.assert_allclose(out.numpy(), expected, rtol=1e-5)
|
|
|
|
def test_scalar_output_stop_gradient(self):
|
|
"""Verify nansum returns scalar for full reduce."""
|
|
for place in self.places:
|
|
paddle.device.set_device(str(place))
|
|
x = paddle.to_tensor([float('nan'), 1.0, 2.0], place=place)
|
|
out = paddle.nansum(x)
|
|
self.assertEqual(out.shape, [])
|
|
|
|
|
|
class TestNansumStaticGraph(unittest.TestCase):
|
|
"""Test nansum under jit.to_static to trigger InferSymbolicShape."""
|
|
|
|
def test_to_static(self):
|
|
class NansumLayer(paddle.nn.Layer):
|
|
def __init__(self):
|
|
super().__init__()
|
|
|
|
def forward(self, x):
|
|
return paddle.nansum(x, axis=1, keepdim=True)
|
|
|
|
net = NansumLayer()
|
|
x_spec = paddle.static.InputSpec(
|
|
shape=[None, None, None], dtype='float32'
|
|
)
|
|
static_net = paddle.jit.to_static(
|
|
net, input_spec=[x_spec], full_graph=True
|
|
)
|
|
x = paddle.randn([2, 3, 4])
|
|
out = static_net(x)
|
|
expected = paddle.nansum(x, axis=1, keepdim=True)
|
|
np.testing.assert_allclose(out.numpy(), expected.numpy())
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|