wehub-resource-sync
339 lines
10 KiB
Python
339 lines
10 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.
|
||
|
||
"""
|
||
CPU unit tests for SetValueGradImpl.
|
||
|
||
Covers the four branches inside SetValueGradImpl
|
||
(paddle/phi/kernels/cpu/set_value_grad_kernel.cc):
|
||
|
||
Branch 1 - x_grad:
|
||
out_grad is copied to x_grad, then the slice region is zeroed out.
|
||
|
||
Branch 2 - value_grad (same shape, need_reverse=false):
|
||
When value_grad.dims() == out_dims and step > 0, the gradient is
|
||
extracted directly via stridedSlice.
|
||
|
||
Branch 3 - value_grad (same shape, need_reverse=true): ** previously uncovered **
|
||
When step < 0, reverse_vector is set and the stridedSlice result is
|
||
reversed before being written to value_grad.
|
||
|
||
Branch 4 - value_grad (broadcast/reduce):
|
||
When value_grad.dims() != out_dims, gradient is accumulated across
|
||
multiple broadcast tiles.
|
||
|
||
The test chain used throughout:
|
||
a = x * x
|
||
a[index] = value * value
|
||
loss = a.sum()
|
||
loss.backward()
|
||
|
||
Using v*v (instead of a direct assignment) makes gradients non-trivial and
|
||
easy to verify analytically:
|
||
d(loss)/d(x[i]) = 2*x[i] for i NOT in slice region; 0 inside region
|
||
d(loss)/d(v[i]) = 2*v[i] (each v[i] appears once as v[i]^2)
|
||
"""
|
||
|
||
import unittest
|
||
|
||
import numpy as np
|
||
|
||
import paddle
|
||
|
||
|
||
class TestSetValueGradXGrad(unittest.TestCase):
|
||
"""
|
||
Branch 1: x_grad path.
|
||
|
||
SetValueGradImpl copies out_grad into x_grad, then zeros the slice region.
|
||
Expected: x_grad[i] = 2*x[i] outside the slice; 0 inside.
|
||
"""
|
||
|
||
def setUp(self):
|
||
paddle.disable_static()
|
||
paddle.set_device('cpu')
|
||
|
||
def test_1d_slice_step1(self):
|
||
x_np = np.arange(1.0, 7.0, dtype='float32') # [1,2,3,4,5,6]
|
||
value_np = np.array([10.0, 10.0], dtype='float32')
|
||
|
||
x = paddle.to_tensor(x_np, stop_gradient=False)
|
||
value = paddle.to_tensor(value_np, stop_gradient=False)
|
||
|
||
a = x * x
|
||
a[1:3] = value # slice indices 1,2
|
||
loss = a.sum()
|
||
loss.backward()
|
||
|
||
expected = 2.0 * x_np
|
||
expected[1:3] = 0.0
|
||
|
||
np.testing.assert_allclose(
|
||
x.grad.numpy(),
|
||
expected,
|
||
rtol=1e-6,
|
||
err_msg=f'x_grad mismatch: expected {expected}, got {x.grad.numpy()}',
|
||
)
|
||
|
||
def test_2d_slice(self):
|
||
x_np = np.arange(1.0, 13.0, dtype='float32').reshape(3, 4)
|
||
value_np = np.ones((2, 4), dtype='float32') * 99.0
|
||
|
||
x = paddle.to_tensor(x_np, stop_gradient=False)
|
||
value = paddle.to_tensor(value_np, stop_gradient=False)
|
||
|
||
a = x * x
|
||
a[0:2, :] = value
|
||
loss = a.sum()
|
||
loss.backward()
|
||
|
||
expected = 2.0 * x_np
|
||
expected[0:2, :] = 0.0
|
||
|
||
np.testing.assert_allclose(
|
||
x.grad.numpy(),
|
||
expected,
|
||
rtol=1e-6,
|
||
)
|
||
|
||
|
||
class TestSetValueGradValueGradSameShape(unittest.TestCase):
|
||
"""
|
||
Branch 2: value_grad when value_grad.dims() == out_dims and step > 0.
|
||
|
||
SetValueGradImpl extracts the gradient via a single stridedSlice and assigns
|
||
it directly to value_grad (no accumulation, no reverse).
|
||
Expected: value_grad[i] = 2*v[i].
|
||
"""
|
||
|
||
def setUp(self):
|
||
paddle.disable_static()
|
||
paddle.set_device('cpu')
|
||
|
||
def test_1d_exact_shape(self):
|
||
x_np = np.arange(1.0, 7.0, dtype='float32')
|
||
value_np = np.array([10.0, 20.0, 30.0], dtype='float32')
|
||
|
||
x = paddle.to_tensor(x_np, stop_gradient=False)
|
||
value = paddle.to_tensor(value_np, stop_gradient=False)
|
||
|
||
a = x * x
|
||
a[2:5] = value * value # slice output shape [3] == value shape [3]
|
||
loss = a.sum()
|
||
loss.backward()
|
||
|
||
expected = 2.0 * value_np
|
||
np.testing.assert_allclose(
|
||
value.grad.numpy(),
|
||
expected,
|
||
rtol=1e-6,
|
||
)
|
||
|
||
def test_2d_exact_shape(self):
|
||
x_np = np.arange(1.0, 13.0, dtype='float32').reshape(4, 3)
|
||
value_np = np.arange(100.0, 106.0, dtype='float32').reshape(2, 3)
|
||
|
||
x = paddle.to_tensor(x_np, stop_gradient=False)
|
||
value = paddle.to_tensor(value_np, stop_gradient=False)
|
||
|
||
a = x * x
|
||
a[1:3, :] = (
|
||
value * value
|
||
) # slice output shape [2,3] == value shape [2,3]
|
||
loss = a.sum()
|
||
loss.backward()
|
||
|
||
expected = 2.0 * value_np
|
||
np.testing.assert_allclose(
|
||
value.grad.numpy(),
|
||
expected,
|
||
rtol=1e-5,
|
||
)
|
||
|
||
|
||
class TestSetValueGradNegativeStep(unittest.TestCase):
|
||
"""
|
||
Branch 3: need_reverse = true (step < 0).
|
||
|
||
When any axis has a negative step, reverse_vector is set for that axis and
|
||
SetValueGradImpl reverses the stridedSlice result before writing to
|
||
value_grad. This branch was NOT covered by existing tests.
|
||
|
||
The key invariant: value_grad[i] = 2*v[i] regardless of the direction in
|
||
which v was written into x, because the reverse restores the correspondence
|
||
between value positions and gradient positions.
|
||
"""
|
||
|
||
def setUp(self):
|
||
paddle.disable_static()
|
||
paddle.set_device('cpu')
|
||
|
||
def test_1d_step_neg1(self):
|
||
"""a[5:2:-1] = v => writes a[5],a[4],a[3]; step=-1 triggers need_reverse."""
|
||
x_np = np.arange(1.0, 7.0, dtype='float32') # [1,2,3,4,5,6]
|
||
value_np = np.array([10.0, 20.0, 30.0], dtype='float32')
|
||
|
||
x = paddle.to_tensor(x_np, stop_gradient=False)
|
||
value = paddle.to_tensor(value_np, stop_gradient=False)
|
||
|
||
a = x * x
|
||
a[5:2:-1] = value * value # fills indices 5,4,3
|
||
loss = a.sum()
|
||
loss.backward()
|
||
|
||
# x_grad: 2*x everywhere, 0 at indices {3,4,5}
|
||
expected_x_grad = 2.0 * x_np
|
||
expected_x_grad[3:6] = 0.0
|
||
|
||
# value_grad: 2*v (reverse restores original order)
|
||
expected_value_grad = 2.0 * value_np
|
||
|
||
np.testing.assert_allclose(
|
||
x.grad.numpy(),
|
||
expected_x_grad,
|
||
rtol=1e-6,
|
||
err_msg=f'x_grad: expected {expected_x_grad}, got {x.grad.numpy()}',
|
||
)
|
||
np.testing.assert_allclose(
|
||
value.grad.numpy(),
|
||
expected_value_grad,
|
||
rtol=1e-6,
|
||
err_msg=f'value_grad: expected {expected_value_grad}, got {value.grad.numpy()}',
|
||
)
|
||
|
||
def test_2d_negative_step_axis0(self):
|
||
"""a[3:0:-1, :] = v => fills rows 3,2,1 in reverse; step=-1 on axis 0."""
|
||
x_np = np.arange(1.0, 13.0, dtype='float32').reshape(4, 3)
|
||
value_np = np.arange(10.0, 19.0, dtype='float32').reshape(3, 3)
|
||
|
||
x = paddle.to_tensor(x_np, stop_gradient=False)
|
||
value = paddle.to_tensor(value_np, stop_gradient=False)
|
||
|
||
a = x * x
|
||
a[3:0:-1, :] = value * value # fills rows 3,2,1
|
||
loss = a.sum()
|
||
loss.backward()
|
||
|
||
# x_grad: 2*x for row 0 only; 0 for rows 1,2,3
|
||
expected_x_grad = 2.0 * x_np
|
||
expected_x_grad[1:4, :] = 0.0
|
||
|
||
expected_value_grad = 2.0 * value_np
|
||
|
||
np.testing.assert_allclose(
|
||
x.grad.numpy(),
|
||
expected_x_grad,
|
||
rtol=1e-6,
|
||
)
|
||
np.testing.assert_allclose(
|
||
value.grad.numpy(),
|
||
expected_value_grad,
|
||
rtol=1e-6,
|
||
)
|
||
|
||
def test_1d_step_neg2(self):
|
||
"""a[::-2] = v => step=-2, hits indices 7,5,3,1."""
|
||
x_np = np.arange(1.0, 9.0, dtype='float32') # [1..8], shape [8]
|
||
value_np = np.array([10.0, 20.0, 30.0, 40.0], dtype='float32')
|
||
|
||
x = paddle.to_tensor(x_np, stop_gradient=False)
|
||
value = paddle.to_tensor(value_np, stop_gradient=False)
|
||
|
||
a = x * x
|
||
a[::-2] = value * value # fills indices 7,5,3,1
|
||
loss = a.sum()
|
||
loss.backward()
|
||
|
||
# x_grad=0 at {1,3,5,7}; 2*x elsewhere
|
||
expected_x_grad = 2.0 * x_np
|
||
expected_x_grad[[1, 3, 5, 7]] = 0.0
|
||
|
||
expected_value_grad = 2.0 * value_np
|
||
|
||
np.testing.assert_allclose(
|
||
x.grad.numpy(),
|
||
expected_x_grad,
|
||
rtol=1e-6,
|
||
)
|
||
np.testing.assert_allclose(
|
||
value.grad.numpy(),
|
||
expected_value_grad,
|
||
rtol=1e-6,
|
||
)
|
||
|
||
|
||
class TestSetValueGradValueGradBroadcast(unittest.TestCase):
|
||
"""
|
||
Branch 4: value_grad when value_grad.dims() != out_dims (broadcast/reduce).
|
||
|
||
SetValueGradImpl accumulates the gradient from multiple broadcast tiles into
|
||
value_grad via repeated additions.
|
||
Expected: value_grad[i] = 2*v[i] * (number of tiles v[i] is broadcast into).
|
||
"""
|
||
|
||
def setUp(self):
|
||
paddle.disable_static()
|
||
paddle.set_device('cpu')
|
||
|
||
def test_scalar_broadcast_into_2d_slice(self):
|
||
"""value shape [1] broadcast into slice shape [2,4] (8 elements)."""
|
||
x_np = np.arange(1.0, 13.0, dtype='float32').reshape(3, 4)
|
||
value_np = np.array([5.0], dtype='float32')
|
||
|
||
x = paddle.to_tensor(x_np, stop_gradient=False)
|
||
value = paddle.to_tensor(value_np, stop_gradient=False)
|
||
|
||
a = x * x
|
||
a[0:2, :] = value * value # slice [2,4], value [1] -> 8 tiles
|
||
loss = a.sum()
|
||
loss.backward()
|
||
|
||
# d(loss)/d(v[0]) = sum over 8 positions of 2*v[0] = 2*5*8 = 80
|
||
expected_value_grad = np.array([2.0 * 5.0 * 8.0], dtype='float32')
|
||
|
||
np.testing.assert_allclose(
|
||
value.grad.numpy(),
|
||
expected_value_grad,
|
||
rtol=1e-5,
|
||
err_msg=f'Expected {expected_value_grad}, got {value.grad.numpy()}',
|
||
)
|
||
|
||
def test_row_vector_broadcast(self):
|
||
"""value shape [1,4] broadcast into slice shape [2,4] (2 rows)."""
|
||
x_np = np.arange(1.0, 13.0, dtype='float32').reshape(3, 4)
|
||
value_np = np.array([[10.0, 20.0, 30.0, 40.0]], dtype='float32')
|
||
|
||
x = paddle.to_tensor(x_np, stop_gradient=False)
|
||
value = paddle.to_tensor(value_np, stop_gradient=False)
|
||
|
||
a = x * x
|
||
a[0:2, :] = (
|
||
value * value
|
||
) # slice [2,4], value [1,4] -> each v[j] used 2×
|
||
loss = a.sum()
|
||
loss.backward()
|
||
|
||
# d(loss)/d(v[0,j]) = 2*v[0,j] * 2 (2 rows use the same v[0,j])
|
||
expected_value_grad = 2.0 * value_np * 2.0
|
||
|
||
np.testing.assert_allclose(
|
||
value.grad.numpy(),
|
||
expected_value_grad,
|
||
rtol=1e-5,
|
||
)
|
||
|
||
|
||
if __name__ == '__main__':
|
||
unittest.main()
|