wehub-resource-sync
735 lines
26 KiB
Python
735 lines
26 KiB
Python
# Copyright (c) 2023 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.
|
|
|
|
import unittest
|
|
from itertools import combinations
|
|
|
|
import numpy as np
|
|
|
|
import paddle
|
|
from paddle.base import Program
|
|
|
|
paddle.enable_static()
|
|
|
|
|
|
def get_all_devices():
|
|
"""Return both cpu and gpu devices for coverage."""
|
|
devices = ['cpu']
|
|
if paddle.is_compiled_with_cuda():
|
|
devices.append('gpu')
|
|
return devices
|
|
|
|
|
|
def compute_index_fill_ref(x, axis, index, value):
|
|
perm = list(range(len(x.shape)))
|
|
perm[0] = axis
|
|
perm[axis] = 0
|
|
|
|
out = np.transpose(x, perm)
|
|
out[index] = value
|
|
out = np.transpose(out, perm)
|
|
return out
|
|
|
|
|
|
def compute_index_fill_grad_ref(x_shape, axis, index):
|
|
grad = np.ones(x_shape, dtype='float64')
|
|
perm = list(range(len(x_shape)))
|
|
perm[0] = axis
|
|
perm[axis] = 0
|
|
grad = np.transpose(grad, perm)
|
|
grad[index] = 0
|
|
grad = np.transpose(grad, perm)
|
|
return grad
|
|
|
|
|
|
class TestIndexFillAPIBase(unittest.TestCase):
|
|
def setUp(self):
|
|
self.init_setting()
|
|
self.modify_setting()
|
|
self.x_np = np.random.random(self.x_shape).astype(self.dtype_np)
|
|
self.index_np = np.array(self.combs[np.random.randint(0, 252)]).astype(
|
|
self.index_type
|
|
)
|
|
|
|
self.place = get_all_devices()
|
|
if self.dtype_np == 'float16' and 'cpu' in self.place:
|
|
self.place.remove('cpu')
|
|
|
|
def init_setting(self):
|
|
self.dtype_np = 'float64'
|
|
self.index_type = 'int64'
|
|
self.x_shape = (20, 40)
|
|
self.index_size = (5,)
|
|
self.axis = 0
|
|
self.value = -1
|
|
self.combs = list(combinations(list(range(10)), self.index_size[0]))
|
|
|
|
def modify_setting(self):
|
|
pass
|
|
|
|
def test_static_graph(self):
|
|
paddle.enable_static()
|
|
for place in self.place:
|
|
with paddle.static.program_guard(Program()):
|
|
x = paddle.static.data(
|
|
name="x", shape=self.x_shape, dtype=self.dtype_np
|
|
)
|
|
index = paddle.static.data(
|
|
name="index", shape=self.index_size, dtype=self.index_type
|
|
)
|
|
out = paddle.index_fill(x, index, self.axis, self.value)
|
|
exe = paddle.static.Executor(place=place)
|
|
feed_list = {"x": self.x_np, "index": self.index_np}
|
|
pd_res = exe.run(
|
|
paddle.static.default_main_program(),
|
|
feed=feed_list,
|
|
fetch_list=[out],
|
|
)[0]
|
|
ref_res = compute_index_fill_ref(
|
|
self.x_np, self.axis, self.index_np, self.value
|
|
)
|
|
np.testing.assert_allclose(ref_res, pd_res)
|
|
|
|
def test_dygraph(self):
|
|
paddle.disable_static()
|
|
for place in self.place:
|
|
paddle.device.set_device(place)
|
|
x_pd = paddle.to_tensor(self.x_np)
|
|
index_pd = paddle.to_tensor(self.index_np)
|
|
pd_res = paddle.index_fill(x_pd, index_pd, self.axis, self.value)
|
|
ref_res = compute_index_fill_ref(
|
|
self.x_np, self.axis, self.index_np, self.value
|
|
)
|
|
np.testing.assert_allclose(ref_res, pd_res)
|
|
|
|
def test_errors(self):
|
|
data_np = np.random.random((10, 10)).astype(np.float32)
|
|
index = paddle.to_tensor([0, 2])
|
|
|
|
def test_index_not_tensor():
|
|
res = paddle.index_fill(data_np, [0, 2], axis=-1, value=-1)
|
|
|
|
self.assertRaises(ValueError, test_index_not_tensor)
|
|
|
|
def test_value_shape():
|
|
res = paddle.index_fill(
|
|
data_np, index, axis=-1, value=paddle.to_tensor([-1, -4])
|
|
)
|
|
|
|
self.assertRaises(ValueError, test_value_shape)
|
|
|
|
def test_axis_range():
|
|
res = paddle.index_fill(data_np, index, axis=4, value=-1)
|
|
|
|
self.assertRaises(ValueError, test_axis_range)
|
|
|
|
|
|
class TestIndexFillAPI1(TestIndexFillAPIBase):
|
|
def modify_setting(self):
|
|
self.dtype_np = 'int64'
|
|
self.index_type = 'int32'
|
|
self.x_shape = (10, 15, 10)
|
|
self.axis = 1
|
|
|
|
|
|
class TestIndexFillAPI2(TestIndexFillAPIBase):
|
|
def modify_setting(self):
|
|
self.dtype_np = 'bool'
|
|
self.index_type = 'int32'
|
|
self.x_shape = (10, 15, 10)
|
|
self.axis = 1
|
|
self.value = True
|
|
|
|
|
|
class TestIndexFillAPI3(TestIndexFillAPIBase):
|
|
def modify_setting(self):
|
|
self.dtype_np = 'float16'
|
|
self.x_shape = (10, 15, 10)
|
|
self.axis = 1
|
|
self.value = 0.5
|
|
|
|
|
|
class TestIndexFillAPINegativeAxis(TestIndexFillAPIBase):
|
|
"""Test negative axis to cover `dim < 0` branch in kernel."""
|
|
|
|
def modify_setting(self):
|
|
self.dtype_np = 'float64'
|
|
self.x_shape = (10, 15, 10)
|
|
self.axis = -1
|
|
self.value = -2.0
|
|
|
|
|
|
class TestIndexFillAPINegativeAxis2(TestIndexFillAPIBase):
|
|
"""Test negative axis=-2 on 3D tensor."""
|
|
|
|
def modify_setting(self):
|
|
self.dtype_np = 'float32'
|
|
self.x_shape = (10, 15, 10)
|
|
self.axis = -2
|
|
self.value = 3.0
|
|
|
|
|
|
class TestIndexFillAPI_ZeroSize(unittest.TestCase):
|
|
def setUp(self):
|
|
self.init_setting()
|
|
self.x_np = np.random.random(self.x_shape).astype(self.dtype_np)
|
|
self.index_np = np.random.random(self.index_size).astype(
|
|
self.index_type
|
|
)
|
|
|
|
self.place = get_all_devices()
|
|
if self.dtype_np == 'float16' and 'cpu' in self.place:
|
|
self.place.remove('cpu')
|
|
|
|
def init_setting(self):
|
|
self.dtype_np = 'float64'
|
|
self.index_type = 'int64'
|
|
self.x_shape = (20, 40)
|
|
# test index with zero size
|
|
self.index_size = (0,)
|
|
self.axis = 0
|
|
self.value = -1
|
|
|
|
def test_dygraph(self):
|
|
paddle.disable_static()
|
|
for place in self.place:
|
|
paddle.device.set_device(place)
|
|
x_pd = paddle.to_tensor(self.x_np)
|
|
x_pd.stop_gradient = False
|
|
index_pd = paddle.to_tensor(self.index_np)
|
|
pd_res = paddle.index_fill(x_pd, index_pd, self.axis, self.value)
|
|
ref_res = compute_index_fill_ref(
|
|
self.x_np, self.axis, self.index_np, self.value
|
|
)
|
|
np.testing.assert_allclose(ref_res, pd_res)
|
|
pd_res.sum().backward()
|
|
np.testing.assert_allclose(
|
|
x_pd.grad.numpy(), np.ones_like(self.x_np)
|
|
)
|
|
paddle.enable_static()
|
|
|
|
|
|
class TestIndexFillAPI_ZeroSize2(TestIndexFillAPI_ZeroSize):
|
|
def init_setting(self):
|
|
self.dtype_np = 'float64'
|
|
self.index_type = 'int64'
|
|
self.x_shape = (20, 0)
|
|
self.index_size = (2,)
|
|
self.axis = 0
|
|
self.value = -1
|
|
|
|
|
|
class TestIndexFillGradBase(unittest.TestCase):
|
|
"""Test backward to cover IndexFillGradKernel on CPU and GPU."""
|
|
|
|
def setUp(self):
|
|
self.init_setting()
|
|
self.modify_setting()
|
|
self.place = get_all_devices()
|
|
|
|
def init_setting(self):
|
|
self.dtype_np = 'float64'
|
|
self.index_type = 'int64'
|
|
self.x_shape = (6, 8)
|
|
self.index_np = np.array([0, 2, 4], dtype=self.index_type)
|
|
self.axis = 0
|
|
self.value = -1.0
|
|
|
|
def modify_setting(self):
|
|
pass
|
|
|
|
def test_dygraph_grad(self):
|
|
paddle.disable_static()
|
|
for place in self.place:
|
|
paddle.device.set_device(place)
|
|
x_np = np.random.random(self.x_shape).astype(self.dtype_np)
|
|
x_pd = paddle.to_tensor(x_np)
|
|
x_pd.stop_gradient = False
|
|
index_pd = paddle.to_tensor(self.index_np)
|
|
out = paddle.index_fill(x_pd, index_pd, self.axis, self.value)
|
|
loss = out.sum()
|
|
loss.backward()
|
|
|
|
expected_grad = compute_index_fill_grad_ref(
|
|
self.x_shape, self.axis, self.index_np
|
|
)
|
|
np.testing.assert_allclose(
|
|
x_pd.grad.numpy(), expected_grad, rtol=1e-5, atol=1e-5
|
|
)
|
|
paddle.enable_static()
|
|
|
|
def test_static_grad(self):
|
|
paddle.enable_static()
|
|
for place in self.place:
|
|
with paddle.static.program_guard(Program()):
|
|
x = paddle.static.data(
|
|
name="x", shape=self.x_shape, dtype=self.dtype_np
|
|
)
|
|
x.stop_gradient = False
|
|
index = paddle.static.data(
|
|
name="index",
|
|
shape=self.index_np.shape,
|
|
dtype=self.index_type,
|
|
)
|
|
out = paddle.index_fill(x, index, self.axis, self.value)
|
|
loss = paddle.mean(out)
|
|
grads = paddle.static.gradients(loss, [x])
|
|
|
|
exe = paddle.static.Executor(place=place)
|
|
x_np = np.random.random(self.x_shape).astype(self.dtype_np)
|
|
res = exe.run(
|
|
paddle.static.default_main_program(),
|
|
feed={"x": x_np, "index": self.index_np},
|
|
fetch_list=[grads[0]],
|
|
)
|
|
self.assertEqual(res[0].shape, tuple(self.x_shape))
|
|
|
|
|
|
class TestIndexFillGradAxis1(TestIndexFillGradBase):
|
|
"""Test grad with axis=1 to cover inner_size > 1 branch."""
|
|
|
|
def modify_setting(self):
|
|
self.x_shape = (4, 6, 5)
|
|
self.index_np = np.array([1, 3], dtype=self.index_type)
|
|
self.axis = 1
|
|
|
|
|
|
class TestIndexFillGradAxis2(TestIndexFillGradBase):
|
|
"""Test grad with axis=2 (last dim) to cover outer_size > 1 branch."""
|
|
|
|
def modify_setting(self):
|
|
self.x_shape = (3, 4, 5)
|
|
self.index_np = np.array([0, 2, 4], dtype=self.index_type)
|
|
self.axis = 2
|
|
|
|
|
|
class TestIndexFillGradNegativeAxis(TestIndexFillGradBase):
|
|
"""Test grad with negative axis to cover dim < 0 branch in grad kernel."""
|
|
|
|
def modify_setting(self):
|
|
self.x_shape = (4, 6, 5)
|
|
self.index_np = np.array([2, 4], dtype=self.index_type)
|
|
self.axis = -1
|
|
|
|
|
|
class TestIndexFillGradInt32Index(TestIndexFillGradBase):
|
|
"""Test grad with int32 index to cover CastToInt64Kernel branch in grad."""
|
|
|
|
def modify_setting(self):
|
|
self.index_type = 'int32'
|
|
self.index_np = np.array([1, 3], dtype=self.index_type)
|
|
|
|
|
|
class TestIndexFillGrad4D(TestIndexFillGradBase):
|
|
"""Test grad with 4D tensor."""
|
|
|
|
def modify_setting(self):
|
|
self.x_shape = (2, 3, 4, 5)
|
|
self.index_np = np.array([0, 2], dtype=self.index_type)
|
|
self.axis = 2
|
|
|
|
|
|
class TestIndexFillGradFloat32(TestIndexFillGradBase):
|
|
"""Test grad with float32 dtype."""
|
|
|
|
def modify_setting(self):
|
|
self.dtype_np = 'float32'
|
|
self.x_shape = (5, 8)
|
|
self.index_np = np.array([1, 4], dtype=self.index_type)
|
|
self.axis = 0
|
|
|
|
|
|
class TestIndexFillGradEmptyIndex(TestIndexFillGradBase):
|
|
"""Test grad with empty index to cover index.numel() == 0 branch in grad."""
|
|
|
|
def modify_setting(self):
|
|
self.index_np = np.array([], dtype=self.index_type)
|
|
|
|
def test_dygraph_grad(self):
|
|
paddle.disable_static()
|
|
for place in self.place:
|
|
paddle.device.set_device(place)
|
|
x_np = np.random.random(self.x_shape).astype(self.dtype_np)
|
|
x_pd = paddle.to_tensor(x_np)
|
|
x_pd.stop_gradient = False
|
|
index_pd = paddle.to_tensor(self.index_np)
|
|
out = paddle.index_fill(x_pd, index_pd, self.axis, self.value)
|
|
loss = out.sum()
|
|
loss.backward()
|
|
np.testing.assert_allclose(
|
|
x_pd.grad.numpy(),
|
|
np.ones(self.x_shape, dtype=self.dtype_np),
|
|
rtol=1e-5,
|
|
atol=1e-5,
|
|
)
|
|
paddle.enable_static()
|
|
|
|
|
|
class TestIndexFillInplace(unittest.TestCase):
|
|
"""Test inplace index_fill_ to cover the inplace kernel path."""
|
|
|
|
def test_inplace_dygraph(self):
|
|
paddle.disable_static()
|
|
for place in get_all_devices():
|
|
paddle.device.set_device(place)
|
|
x_np = np.random.random((5, 6)).astype('float64')
|
|
x_pd = paddle.to_tensor(x_np)
|
|
index_pd = paddle.to_tensor(np.array([1, 3], dtype='int64'))
|
|
paddle.index_fill_(x_pd, index_pd, axis=0, value=-1.0)
|
|
|
|
expected = x_np.copy()
|
|
expected[[1, 3], :] = -1.0
|
|
np.testing.assert_allclose(x_pd.numpy(), expected)
|
|
paddle.enable_static()
|
|
|
|
def test_inplace_int32_index(self):
|
|
paddle.disable_static()
|
|
for place in get_all_devices():
|
|
paddle.device.set_device(place)
|
|
x_np = np.random.random((4, 5)).astype('float32')
|
|
x_pd = paddle.to_tensor(x_np)
|
|
index_pd = paddle.to_tensor(np.array([0, 2], dtype='int32'))
|
|
paddle.index_fill_(x_pd, index_pd, axis=1, value=0.0)
|
|
|
|
expected = x_np.copy()
|
|
expected[:, [0, 2]] = 0.0
|
|
np.testing.assert_allclose(x_pd.numpy(), expected)
|
|
paddle.enable_static()
|
|
|
|
def test_inplace_negative_axis(self):
|
|
paddle.disable_static()
|
|
for place in get_all_devices():
|
|
paddle.device.set_device(place)
|
|
x_np = np.random.random((3, 4, 5)).astype('float64')
|
|
x_pd = paddle.to_tensor(x_np)
|
|
index_pd = paddle.to_tensor(np.array([1, 3], dtype='int64'))
|
|
paddle.index_fill_(x_pd, index_pd, axis=-1, value=2.0)
|
|
|
|
expected = x_np.copy()
|
|
expected[:, :, [1, 3]] = 2.0
|
|
np.testing.assert_allclose(x_pd.numpy(), expected)
|
|
paddle.enable_static()
|
|
|
|
|
|
class TestIndexFillNegativeIndex(unittest.TestCase):
|
|
"""Test negative index values to cover `idx += dim_size` branch in kernel."""
|
|
|
|
def test_negative_index_forward(self):
|
|
paddle.disable_static()
|
|
for place in get_all_devices():
|
|
paddle.device.set_device(place)
|
|
x_np = np.random.random((5, 6)).astype('float64')
|
|
x_pd = paddle.to_tensor(x_np)
|
|
index_pd = paddle.to_tensor(np.array([-1, -3], dtype='int64'))
|
|
pd_res = paddle.index_fill(x_pd, index_pd, 0, -1.0)
|
|
|
|
expected = x_np.copy()
|
|
expected[[-1, -3], :] = -1.0
|
|
np.testing.assert_allclose(pd_res.numpy(), expected)
|
|
paddle.enable_static()
|
|
|
|
def test_negative_index_backward(self):
|
|
paddle.disable_static()
|
|
for place in get_all_devices():
|
|
paddle.device.set_device(place)
|
|
x_np = np.random.random((4, 6)).astype('float64')
|
|
x_pd = paddle.to_tensor(x_np)
|
|
x_pd.stop_gradient = False
|
|
index_pd = paddle.to_tensor(np.array([-1, -2], dtype='int64'))
|
|
out = paddle.index_fill(x_pd, index_pd, 0, -1.0)
|
|
loss = out.sum()
|
|
loss.backward()
|
|
|
|
expected_grad = np.ones_like(x_np)
|
|
expected_grad[[-1, -2], :] = 0
|
|
np.testing.assert_allclose(
|
|
x_pd.grad.numpy(), expected_grad, rtol=1e-5, atol=1e-5
|
|
)
|
|
paddle.enable_static()
|
|
|
|
|
|
class TestIndexFillErrorCases(unittest.TestCase):
|
|
"""Test error cases to cover PADDLE_ENFORCE branches in InferMeta."""
|
|
|
|
def test_invalid_axis_static(self):
|
|
paddle.enable_static()
|
|
with (
|
|
self.assertRaises(ValueError),
|
|
paddle.static.program_guard(Program()),
|
|
):
|
|
x = paddle.static.data(name="x_err", shape=(5, 6), dtype='float32')
|
|
index = paddle.static.data(
|
|
name="idx_err", shape=(2,), dtype='int64'
|
|
)
|
|
paddle.index_fill(x, index, axis=5, value=-1)
|
|
|
|
def test_index_not_1d_static(self):
|
|
paddle.enable_static()
|
|
with (
|
|
self.assertRaises(ValueError),
|
|
paddle.static.program_guard(Program()),
|
|
):
|
|
x = paddle.static.data(name="x_err2", shape=(5, 6), dtype='float32')
|
|
index = paddle.static.data(
|
|
name="idx_err2", shape=(2, 3), dtype='int64'
|
|
)
|
|
paddle.index_fill(x, index, axis=0, value=-1)
|
|
|
|
|
|
class TestIndexFillPyTorchStyleArgs(unittest.TestCase):
|
|
"""Test PyTorch-style argument order to cover decorator_utils branches."""
|
|
|
|
def test_pytorch_positional_order(self):
|
|
"""Cover: isinstance(args[1], int) branch in index_fill_decorator."""
|
|
paddle.disable_static()
|
|
for place in get_all_devices():
|
|
paddle.device.set_device(place)
|
|
x_np = np.random.random((5, 6)).astype('float64')
|
|
x_pd = paddle.to_tensor(x_np)
|
|
index_pd = paddle.to_tensor(np.array([0, 2], dtype='int64'))
|
|
# PyTorch order: (input, dim, index, value)
|
|
res = paddle.index_fill(x_pd, 0, index_pd, -1.0)
|
|
expected = x_np.copy()
|
|
expected[[0, 2], :] = -1.0
|
|
np.testing.assert_allclose(res.numpy(), expected)
|
|
paddle.enable_static()
|
|
|
|
def test_pytorch_keyword_aliases(self):
|
|
"""Cover: 'input' and 'dim' keyword alias branches."""
|
|
paddle.disable_static()
|
|
for place in get_all_devices():
|
|
paddle.device.set_device(place)
|
|
x_np = np.random.random((4, 5)).astype('float64')
|
|
x_pd = paddle.to_tensor(x_np)
|
|
index_pd = paddle.to_tensor(np.array([1, 3], dtype='int64'))
|
|
res = paddle.index_fill(
|
|
input=x_pd, dim=0, index=index_pd, value=-2.0
|
|
)
|
|
expected = x_np.copy()
|
|
expected[[1, 3], :] = -2.0
|
|
np.testing.assert_allclose(res.numpy(), expected)
|
|
paddle.enable_static()
|
|
|
|
def test_inplace_pytorch_keyword(self):
|
|
"""Cover: 'input'/'dim' keyword alias in index_fill_ inplace."""
|
|
paddle.disable_static()
|
|
for place in get_all_devices():
|
|
paddle.device.set_device(place)
|
|
x_np = np.random.random((4, 5)).astype('float64')
|
|
x_pd = paddle.to_tensor(x_np)
|
|
index_pd = paddle.to_tensor(np.array([0, 2], dtype='int64'))
|
|
paddle.index_fill_(input=x_pd, dim=1, index=index_pd, value=0.0)
|
|
expected = x_np.copy()
|
|
expected[:, [0, 2]] = 0.0
|
|
np.testing.assert_allclose(x_pd.numpy(), expected)
|
|
paddle.enable_static()
|
|
|
|
|
|
class TestIndexFillValueBranches(unittest.TestCase):
|
|
"""Test value-related branches in _index_fill_impl."""
|
|
|
|
def test_value_as_0d_tensor(self):
|
|
"""Cover: value is Variable with numel==1, goes through float(value)."""
|
|
paddle.disable_static()
|
|
for place in get_all_devices():
|
|
paddle.device.set_device(place)
|
|
x_np = np.random.random((4, 5)).astype('float64')
|
|
x_pd = paddle.to_tensor(x_np)
|
|
index_pd = paddle.to_tensor(np.array([1, 3], dtype='int64'))
|
|
value_pd = paddle.to_tensor(-1.0)
|
|
res = paddle.index_fill(x_pd, index_pd, axis=0, value=value_pd)
|
|
expected = x_np.copy()
|
|
expected[[1, 3], :] = -1.0
|
|
np.testing.assert_allclose(res.numpy(), expected)
|
|
paddle.enable_static()
|
|
|
|
def test_value_non_scalar_tensor_error(self):
|
|
"""Cover: value is Variable with len(value.shape) > 0."""
|
|
paddle.disable_static()
|
|
for place in get_all_devices():
|
|
paddle.device.set_device(place)
|
|
x_pd = paddle.to_tensor(np.random.random((4, 5)).astype('float64'))
|
|
index_pd = paddle.to_tensor(np.array([1], dtype='int64'))
|
|
value_pd = paddle.to_tensor([-1.0, -2.0])
|
|
with self.assertRaises(ValueError):
|
|
paddle.index_fill(x_pd, index_pd, axis=0, value=value_pd)
|
|
paddle.enable_static()
|
|
|
|
def test_value_as_int(self):
|
|
"""Cover: value is plain int, goes through float(value) non-Variable branch."""
|
|
paddle.disable_static()
|
|
for place in get_all_devices():
|
|
paddle.device.set_device(place)
|
|
x_np = np.random.random((3, 4)).astype('float64')
|
|
x_pd = paddle.to_tensor(x_np)
|
|
index_pd = paddle.to_tensor(np.array([0, 2], dtype='int64'))
|
|
res = paddle.index_fill(x_pd, index_pd, axis=1, value=0)
|
|
expected = x_np.copy()
|
|
expected[:, [0, 2]] = 0
|
|
np.testing.assert_allclose(res.numpy(), expected)
|
|
paddle.enable_static()
|
|
|
|
|
|
class TestIndexFillGradNegativeAxis2(TestIndexFillGradBase):
|
|
"""Test grad with negative axis=-2 to cover dim < 0 branch in grad kernel on 3D."""
|
|
|
|
def modify_setting(self):
|
|
self.x_shape = (3, 4, 5)
|
|
self.index_np = np.array([1, 2], dtype=self.index_type)
|
|
self.axis = -2
|
|
|
|
|
|
class TestIndexFillGradInt32IndexAxis1(TestIndexFillGradBase):
|
|
"""Test grad with int32 index on axis=1."""
|
|
|
|
def modify_setting(self):
|
|
self.index_type = 'int32'
|
|
self.x_shape = (3, 6, 4)
|
|
self.index_np = np.array([0, 3, 5], dtype=self.index_type)
|
|
self.axis = 1
|
|
|
|
|
|
class TestIndexFillOptimizationCorrectness(unittest.TestCase):
|
|
"""Tests specifically targeting the optimization changes:
|
|
1. IndexT template (int32 vs int64 based on numel)
|
|
2. IndT template (int32/int64 index without CastToInt64)
|
|
3. Python-side value handling (no to_tensor roundtrip)
|
|
"""
|
|
|
|
def _run_and_check(self, x_np, index_np, axis, value, dtype_np='float64'):
|
|
ref = compute_index_fill_ref(x_np, axis, index_np, value)
|
|
for place in get_all_devices():
|
|
paddle.device.set_device(place)
|
|
x_pd = paddle.to_tensor(x_np)
|
|
idx_pd = paddle.to_tensor(index_np)
|
|
res = paddle.index_fill(x_pd, idx_pd, axis, value)
|
|
np.testing.assert_allclose(res.numpy(), ref, rtol=1e-5)
|
|
|
|
def test_int32_index_direct_pass(self):
|
|
"""IndT=int32: index with int32 dtype should work without cast."""
|
|
paddle.disable_static()
|
|
x_np = np.random.random((8, 10)).astype('float64')
|
|
index_np = np.array([1, 3, 5], dtype='int32')
|
|
self._run_and_check(x_np, index_np, axis=0, value=-1.0)
|
|
paddle.enable_static()
|
|
|
|
def test_int64_index_direct_pass(self):
|
|
"""IndT=int64: index with int64 dtype should work directly."""
|
|
paddle.disable_static()
|
|
x_np = np.random.random((8, 10)).astype('float64')
|
|
index_np = np.array([0, 4, 7], dtype='int64')
|
|
self._run_and_check(x_np, index_np, axis=0, value=2.5)
|
|
paddle.enable_static()
|
|
|
|
def test_small_numel_int32_path(self):
|
|
"""IndexT=int32: small tensor numel fits in int32."""
|
|
paddle.disable_static()
|
|
x_np = np.random.random((16, 32, 8)).astype('float64')
|
|
index_np = np.array([0, 5, 10, 15], dtype='int64')
|
|
self._run_and_check(x_np, index_np, axis=1, value=-0.5)
|
|
paddle.enable_static()
|
|
|
|
def test_value_python_float(self):
|
|
"""Python float value should pass directly without to_tensor."""
|
|
paddle.disable_static()
|
|
x_np = np.random.random((4, 6)).astype('float32')
|
|
index_np = np.array([1, 3], dtype='int64')
|
|
self._run_and_check(x_np, index_np, axis=1, value=3.14)
|
|
paddle.enable_static()
|
|
|
|
def test_value_python_int(self):
|
|
"""Python int value should pass directly."""
|
|
paddle.disable_static()
|
|
x_np = np.random.randint(0, 100, (5, 8)).astype('int64')
|
|
index_np = np.array([0, 2, 4], dtype='int64')
|
|
self._run_and_check(x_np, index_np, axis=0, value=999)
|
|
paddle.enable_static()
|
|
|
|
def test_value_python_bool(self):
|
|
"""Python bool value should work with bool tensor."""
|
|
paddle.disable_static()
|
|
x_np = np.random.choice([True, False], size=(4, 6)).astype('bool')
|
|
index_np = np.array([1, 3], dtype='int64')
|
|
self._run_and_check(x_np, index_np, axis=0, value=True)
|
|
paddle.enable_static()
|
|
|
|
def test_value_zero_dim_tensor(self):
|
|
"""0-D Tensor value should be extracted via .item()."""
|
|
paddle.disable_static()
|
|
for place in get_all_devices():
|
|
paddle.device.set_device(place)
|
|
x_np = np.random.random((6, 8)).astype('float64')
|
|
index_np = np.array([2, 5], dtype='int64')
|
|
x_pd = paddle.to_tensor(x_np)
|
|
idx_pd = paddle.to_tensor(index_np)
|
|
val_pd = paddle.to_tensor(-7.0)
|
|
res = paddle.index_fill(x_pd, idx_pd, axis=0, value=val_pd)
|
|
ref = compute_index_fill_ref(x_np, 0, index_np, -7.0)
|
|
np.testing.assert_allclose(res.numpy(), ref)
|
|
paddle.enable_static()
|
|
|
|
def test_inplace_with_scalar_value(self):
|
|
"""Inplace variant should also work with scalar value."""
|
|
paddle.disable_static()
|
|
for place in get_all_devices():
|
|
paddle.device.set_device(place)
|
|
x_np = np.random.random((4, 5)).astype('float64')
|
|
index_np = np.array([0, 3], dtype='int32')
|
|
x_pd = paddle.to_tensor(x_np)
|
|
idx_pd = paddle.to_tensor(index_np)
|
|
paddle.index_fill_(x_pd, idx_pd, axis=1, value=-2.0)
|
|
ref = compute_index_fill_ref(x_np, 1, index_np, -2.0)
|
|
np.testing.assert_allclose(x_pd.numpy(), ref)
|
|
paddle.enable_static()
|
|
|
|
def test_int32_index_with_various_dtypes(self):
|
|
"""int32 index with float16/float32/float64/int32/int64 tensors."""
|
|
paddle.disable_static()
|
|
index_np = np.array([0, 2], dtype='int32')
|
|
for dtype in ['float32', 'float64', 'int32', 'int64']:
|
|
if dtype.startswith('int'):
|
|
x_np = np.random.randint(0, 50, (5, 4)).astype(dtype)
|
|
else:
|
|
x_np = np.random.random((5, 4)).astype(dtype)
|
|
self._run_and_check(x_np, index_np, axis=0, value=-1)
|
|
paddle.enable_static()
|
|
|
|
def test_negative_axis(self):
|
|
"""Negative axis with optimized kernel."""
|
|
paddle.disable_static()
|
|
x_np = np.random.random((3, 4, 5)).astype('float64')
|
|
index_np = np.array([1, 3], dtype='int64')
|
|
self._run_and_check(x_np, index_np, axis=-1, value=0.0)
|
|
paddle.enable_static()
|
|
|
|
def test_single_index(self):
|
|
"""Single index element - minimal workload."""
|
|
paddle.disable_static()
|
|
x_np = np.random.random((10, 20)).astype('float64')
|
|
index_np = np.array([5], dtype='int32')
|
|
self._run_and_check(x_np, index_np, axis=0, value=42.0)
|
|
paddle.enable_static()
|
|
|
|
def test_full_dim_index(self):
|
|
"""Index covers entire dimension."""
|
|
paddle.disable_static()
|
|
x_np = np.random.random((4, 6)).astype('float64')
|
|
index_np = np.arange(6, dtype='int64')
|
|
self._run_and_check(x_np, index_np, axis=1, value=-1.0)
|
|
paddle.enable_static()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|