1005 lines
28 KiB
Python
1005 lines
28 KiB
Python
# Copyright (c) 2018 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
|
|
|
|
import numpy as np
|
|
from op_test import (
|
|
OpTest,
|
|
check_cudnn_version_and_compute_capability,
|
|
convert_float_to_uint16,
|
|
get_device_place,
|
|
get_places,
|
|
is_custom_device,
|
|
)
|
|
from utils import dygraph_guard
|
|
|
|
import paddle
|
|
from paddle import base
|
|
from paddle.base.dygraph.base import switch_to_static_graph
|
|
from paddle.framework import core
|
|
|
|
|
|
def gather_numpy(x, index, axis):
|
|
x_transpose = np.swapaxes(x, 0, axis)
|
|
tmp_gather = x_transpose[index, ...]
|
|
gather = np.swapaxes(tmp_gather, 0, axis)
|
|
return gather
|
|
|
|
|
|
class TestGatherOp(OpTest):
|
|
def setUp(self):
|
|
self.op_type = "gather"
|
|
self.python_api = paddle.gather
|
|
self.public_python_api = paddle.gather
|
|
self.config()
|
|
self.prim_op_type = "prim"
|
|
self.init_inputs_and_outputs()
|
|
self.if_enable_cinn()
|
|
|
|
def test_check_output(self):
|
|
self.check_output(check_pir=True, check_symbol_infer=False)
|
|
|
|
def test_check_grad(self):
|
|
self.check_grad(['X'], 'Out', check_pir=True, check_prim_pir=True)
|
|
|
|
def config(self):
|
|
"""
|
|
For multi-dimension input
|
|
"""
|
|
self.x_shape = (10, 20)
|
|
self.config_dtype()
|
|
self.index = [1, 3, 5]
|
|
self.index_type = "int32"
|
|
|
|
def config_dtype(self):
|
|
self.x_type = "float64"
|
|
|
|
def init_inputs_and_outputs(self):
|
|
xnp = np.random.random(self.x_shape).astype(self.x_type)
|
|
if self.x_type == 'complex64' or self.x_type == "cpmolex128":
|
|
xnp = (
|
|
np.random.randint(-10, 10, size=(10, 10))
|
|
+ 1j * np.random.randint(-10, 10, size=(10, 10))
|
|
).astype(self.x_type)
|
|
self.inputs = {
|
|
'X': xnp,
|
|
'Index': np.array(self.index).astype(self.index_type),
|
|
}
|
|
self.outputs = {'Out': self.inputs["X"][self.inputs["Index"]]}
|
|
|
|
def if_enable_cinn(self):
|
|
pass
|
|
|
|
|
|
class TestGatherOp_ZeroDim(TestGatherOp):
|
|
def config(self):
|
|
"""
|
|
For multi-dimension input
|
|
"""
|
|
self.x_shape = 100
|
|
self.config_dtype()
|
|
self.index = 2
|
|
self.index_type = "int32"
|
|
|
|
def if_enable_cinn(self):
|
|
self.enable_cinn = False
|
|
|
|
|
|
class TestGatherOpFP16(TestGatherOp):
|
|
def config_dtype(self):
|
|
self.x_type = "float16"
|
|
|
|
|
|
@unittest.skipIf(
|
|
not check_cudnn_version_and_compute_capability(8100, 8),
|
|
"only support compiled with CUDA or custom device, and for CUDA cudnn version need larger than 8.1.0 and device's compute capability is at least 8.0",
|
|
)
|
|
class TestGatherOpBFP16(TestGatherOp):
|
|
def config_dtype(self):
|
|
self.x_type = "float32"
|
|
self.dtype = np.uint16
|
|
|
|
def init_inputs_and_outputs(self):
|
|
xnp = np.random.random(self.x_shape).astype(self.x_type)
|
|
self.inputs = {
|
|
'X': convert_float_to_uint16(xnp),
|
|
'Index': np.array(self.index).astype(self.index_type),
|
|
}
|
|
self.outputs = {
|
|
'Out': convert_float_to_uint16(xnp[self.inputs["Index"]])
|
|
}
|
|
|
|
def if_enable_cinn(self):
|
|
self.enable_cinn = False
|
|
|
|
def test_check_output(self):
|
|
self.check_output_with_place(
|
|
place=get_device_place(), check_pir=True, check_symbol_infer=False
|
|
)
|
|
|
|
def test_check_grad(self):
|
|
self.check_grad_with_place(
|
|
get_device_place(),
|
|
['X'],
|
|
'Out',
|
|
check_pir=True,
|
|
check_prim_pir=True,
|
|
)
|
|
|
|
|
|
class TestGatherOpComplex64(TestGatherOp):
|
|
def config_dtype(self):
|
|
self.x_type = "complex64"
|
|
|
|
def test_check_grad(self):
|
|
self.check_grad(['X'], 'Out')
|
|
|
|
|
|
class TestGatherOpComplex128(TestGatherOp):
|
|
def config_dtype(self):
|
|
self.x_type = "complex128"
|
|
|
|
def test_check_grad(self):
|
|
self.check_grad(['X'], 'Out')
|
|
|
|
|
|
class TestCase1(TestGatherOp):
|
|
def config(self):
|
|
"""
|
|
For one dimension input
|
|
"""
|
|
self.x_shape = 100
|
|
self.config_dtype()
|
|
self.index = [1, 3, 5]
|
|
self.index_type = "int32"
|
|
|
|
def config_dtype(self):
|
|
self.x_type = "float64"
|
|
|
|
|
|
class TestCase1FP16(TestCase1):
|
|
def config_dtype(self):
|
|
self.x_type = "float16"
|
|
|
|
|
|
class TestCase1BFP16(TestGatherOpBFP16):
|
|
def config(self):
|
|
self.x_shape = 100
|
|
self.config_dtype()
|
|
self.index = [1, 3, 5]
|
|
self.index_type = "int32"
|
|
|
|
|
|
class TestCase1Complex64(TestCase1):
|
|
def config_dtype(self):
|
|
self.x_type = "complex64"
|
|
|
|
def test_check_grad(self):
|
|
self.check_grad(['X'], 'Out')
|
|
|
|
|
|
class TestCase1Complex128(TestCase1):
|
|
def config_dtype(self):
|
|
self.x_type = "complex128"
|
|
|
|
def test_check_grad(self):
|
|
self.check_grad(['X'], 'Out')
|
|
|
|
|
|
class TestCase2(TestGatherOp):
|
|
def config(self):
|
|
"""
|
|
For int64_t index type
|
|
"""
|
|
self.x_shape = 100
|
|
self.config_dtype()
|
|
self.index = [1, 3, 5]
|
|
self.index_type = "int64"
|
|
|
|
def config_dtype(self):
|
|
self.x_type = "float64"
|
|
|
|
|
|
class TestCase2FP16(TestCase2):
|
|
def config_dtype(self):
|
|
self.x_type = "float16"
|
|
|
|
|
|
class TestCase2BFP16(TestGatherOpBFP16):
|
|
def config(self):
|
|
self.x_shape = 100
|
|
self.config_dtype()
|
|
self.index = [1, 3, 5]
|
|
self.index_type = "int64"
|
|
|
|
|
|
class TestCase2Complex64(TestCase2):
|
|
def config_dtype(self):
|
|
self.x_type = "complex64"
|
|
|
|
def test_check_grad(self):
|
|
self.check_grad(['X'], 'Out')
|
|
|
|
|
|
class TestCase2Complex128(TestCase2):
|
|
def config_dtype(self):
|
|
self.x_type = "complex128"
|
|
|
|
def test_check_grad(self):
|
|
self.check_grad(['X'], 'Out')
|
|
|
|
|
|
class TestCase3(TestGatherOp):
|
|
def config(self):
|
|
"""
|
|
For other input type
|
|
"""
|
|
self.x_shape = (10, 20)
|
|
self.config_dtype()
|
|
self.index = [1, 3, 5]
|
|
self.index_type = "int64"
|
|
|
|
def config_dtype(self):
|
|
self.x_type = "float64"
|
|
|
|
|
|
class TestCase3Fp16(TestCase3):
|
|
def config_dtype(self):
|
|
self.x_type = "float16"
|
|
|
|
|
|
class TestCase3BFP16(TestGatherOpBFP16):
|
|
def config(self):
|
|
self.x_shape = (10, 20)
|
|
self.config_dtype()
|
|
self.index = [1, 3, 5]
|
|
self.index_type = "int64"
|
|
|
|
|
|
class TestCase3Complex64(TestCase3):
|
|
def config_dtype(self):
|
|
self.x_type = "complex64"
|
|
|
|
def test_check_grad(self):
|
|
self.check_grad(['X'], 'Out')
|
|
|
|
|
|
class TestCase3Complex128(TestCase3):
|
|
def config_dtype(self):
|
|
self.x_type = "complex128"
|
|
|
|
def test_check_grad(self):
|
|
self.check_grad(['X'], 'Out')
|
|
|
|
|
|
class TestCase4(TestGatherOp):
|
|
def config(self):
|
|
self.x_shape = (10, 20)
|
|
self.attrs = {'overwrite': False}
|
|
self.config_dtype()
|
|
self.index = [1, 1]
|
|
self.index_type = "int32"
|
|
|
|
def config_dtype(self):
|
|
self.x_type = "float64"
|
|
|
|
|
|
class TestCase4FP16(TestCase4):
|
|
def config_dtype(self):
|
|
self.x_type = "float16"
|
|
|
|
|
|
class TestCase4BFP16(TestGatherOpBFP16):
|
|
def config(self):
|
|
self.x_shape = (10, 20)
|
|
self.attrs = {'overwrite': False}
|
|
self.config_dtype()
|
|
self.index = [1, 1]
|
|
self.index_type = "int32"
|
|
|
|
|
|
class TestCase4Complex64(TestCase4):
|
|
def config_dtype(self):
|
|
self.x_type = "complex64"
|
|
|
|
def test_check_grad(self):
|
|
self.check_grad(['X'], 'Out')
|
|
|
|
|
|
class TestCase4Complex128(TestCase4):
|
|
def config_dtype(self):
|
|
self.x_type = "complex128"
|
|
|
|
def test_check_grad(self):
|
|
self.check_grad(['X'], 'Out')
|
|
|
|
|
|
class TestCase5(TestGatherOp):
|
|
def config(self):
|
|
self.x_shape = (10, 20)
|
|
self.attrs = {'overwrite': False}
|
|
self.config_dtype()
|
|
self.index = [1, 1, 3]
|
|
self.index_type = "int32"
|
|
|
|
def config_dtype(self):
|
|
self.x_type = "float64"
|
|
|
|
|
|
class TestCase5BFP16(TestGatherOpBFP16):
|
|
def config(self):
|
|
self.x_shape = (10, 20)
|
|
self.attrs = {'overwrite': False}
|
|
self.config_dtype()
|
|
self.index = [1, 1]
|
|
self.index_type = "int32"
|
|
|
|
|
|
class TestCase5FP16(TestCase5):
|
|
def config_dtype(self):
|
|
self.x_type = "float16"
|
|
|
|
|
|
class TestCase5Complex64(TestCase5):
|
|
def config_dtype(self):
|
|
self.x_type = "complex64"
|
|
|
|
def test_check_grad(self):
|
|
self.check_grad(['X'], 'Out')
|
|
|
|
|
|
class TestCase5Complex128(TestCase5):
|
|
def config_dtype(self):
|
|
self.x_type = "complex128"
|
|
|
|
def test_check_grad(self):
|
|
self.check_grad(['X'], 'Out')
|
|
|
|
|
|
class TestCase6(TestGatherOp):
|
|
def config(self):
|
|
self.x_shape = (10, 20)
|
|
self.attrs = {'overwrite': True}
|
|
self.config_dtype()
|
|
self.index = [1, 3]
|
|
self.index_type = "int32"
|
|
|
|
def config_dtype(self):
|
|
self.x_type = "float64"
|
|
|
|
|
|
class TestCase6FP16(TestCase6):
|
|
def config_dtype(self):
|
|
self.x_type = "float16"
|
|
|
|
|
|
class TestCase6BFP16(TestGatherOpBFP16):
|
|
def config(self):
|
|
self.x_shape = (10, 20)
|
|
self.attrs = {'overwrite': True}
|
|
self.config_dtype()
|
|
self.index = [1, 3]
|
|
self.index_type = "int32"
|
|
|
|
|
|
class TestGatherBF16Op(OpTest):
|
|
def setUp(self):
|
|
self.op_type = "gather"
|
|
self.python_api = paddle.gather
|
|
self.dtype = np.uint16
|
|
self.config()
|
|
xnp = np.random.random(self.x_shape).astype(np.float32)
|
|
axis_np = np.array(self.axis).astype(self.axis_type)
|
|
index_np = np.array(self.index).astype(self.index_type)
|
|
self.inputs = {
|
|
'X': convert_float_to_uint16(xnp),
|
|
'Index': index_np,
|
|
'Axis': axis_np,
|
|
}
|
|
out = gather_numpy(self.inputs['X'], index_np, axis_np[0])
|
|
self.outputs = {'Out': out}
|
|
|
|
def test_check_output(self):
|
|
self.check_output(check_pir=True, check_symbol_infer=False)
|
|
|
|
def test_check_grad(self):
|
|
self.check_grad(['X'], 'Out', numeric_grad_delta=0.5, check_pir=True)
|
|
|
|
def config(self):
|
|
"""
|
|
For multi-dimension input
|
|
"""
|
|
self.x_shape = (3, 88, 3)
|
|
self.index = [1, 3, 5]
|
|
self.index_type = "int32"
|
|
self.axis = [1]
|
|
self.axis_type = "int32"
|
|
|
|
|
|
class TestGatherNegativeAxis(OpTest):
|
|
def setUp(self):
|
|
self.op_type = "gather"
|
|
self.python_api = paddle.gather
|
|
self.dtype = np.uint16
|
|
self.config()
|
|
xnp = np.random.random(self.x_shape).astype(np.float32)
|
|
axis_np = np.array(self.axis).astype(self.axis_type)
|
|
index_np = np.array(self.index).astype(self.index_type)
|
|
self.inputs = {
|
|
'X': convert_float_to_uint16(xnp),
|
|
'Index': index_np,
|
|
'Axis': axis_np,
|
|
}
|
|
out = gather_numpy(self.inputs['X'], index_np, axis_np[0])
|
|
self.outputs = {'Out': out}
|
|
|
|
def test_check_output(self):
|
|
places = [paddle.CPUPlace()]
|
|
if core.is_compiled_with_cuda() or is_custom_device():
|
|
places.append(get_device_place())
|
|
for place in places:
|
|
self.check_output_with_place(place)
|
|
|
|
def test_check_grad(self):
|
|
places = [paddle.CPUPlace()]
|
|
if core.is_compiled_with_cuda() or is_custom_device():
|
|
places.append(get_device_place())
|
|
for place in places:
|
|
self.check_grad_with_place(
|
|
place, ['X'], 'Out', numeric_grad_delta=0.5
|
|
)
|
|
|
|
def config(self):
|
|
"""
|
|
For multi-dimension input
|
|
"""
|
|
self.x_shape = (100, 3)
|
|
self.index = [0, 1, -2]
|
|
self.index_type = "int32"
|
|
self.axis = [-1]
|
|
self.axis_type = "int32"
|
|
|
|
|
|
class TestOutOfRangeError(unittest.TestCase):
|
|
def test_dygraph_forward_and_backward(self):
|
|
with dygraph_guard():
|
|
x = paddle.randn([100, 3]).cpu()
|
|
x.stop_gradient = False
|
|
y = paddle.gather(
|
|
x,
|
|
paddle.to_tensor([0, -2]).cpu(),
|
|
axis=-1,
|
|
)
|
|
grad_x = paddle.grad(y, x)
|
|
|
|
def test_dygraph_error(self):
|
|
with dygraph_guard():
|
|
# out of lower bound
|
|
with self.assertRaises(IndexError):
|
|
_ = paddle.gather(
|
|
paddle.randn([100, 3]).cpu(),
|
|
paddle.to_tensor([0, -4]).cpu(),
|
|
axis=1,
|
|
)
|
|
# out of upper bound
|
|
with self.assertRaises(IndexError):
|
|
_ = paddle.gather(
|
|
paddle.randn([100, 3]).cpu(),
|
|
paddle.to_tensor([0, 3]).cpu(),
|
|
axis=1,
|
|
)
|
|
|
|
|
|
class TestCase6Complex64(TestCase6):
|
|
def config_dtype(self):
|
|
self.x_type = "complex64"
|
|
|
|
def test_check_grad(self):
|
|
self.check_grad(['X'], 'Out')
|
|
|
|
|
|
class TestCase6Complex128(TestCase6):
|
|
def config_dtype(self):
|
|
self.x_type = "complex128"
|
|
|
|
def test_check_grad(self):
|
|
self.check_grad(['X'], 'Out')
|
|
|
|
|
|
class TestGatherOp1(OpTest):
|
|
def setUp(self):
|
|
self.op_type = "gather"
|
|
self.python_api = paddle.gather
|
|
self.config()
|
|
xnp = np.random.random(self.x_shape).astype(self.x_type)
|
|
axis_np = np.array(self.axis).astype(self.index_type)
|
|
index_np = np.array(self.index).astype(self.index_type)
|
|
out = gather_numpy(xnp, index_np, axis_np[0])
|
|
self.inputs = {'X': xnp, 'Index': index_np, 'Axis': axis_np}
|
|
self.outputs = {'Out': out}
|
|
|
|
def test_check_output(self):
|
|
self.check_output(check_pir=True, check_symbol_infer=False)
|
|
|
|
def test_check_grad(self):
|
|
self.check_grad(['X'], 'Out', check_pir=True)
|
|
|
|
def config(self):
|
|
"""
|
|
For multi-dimension input
|
|
"""
|
|
self.x_shape = (3, 88, 3)
|
|
self.config_dtype()
|
|
self.index = [1, 3, 5]
|
|
self.index_type = "int32"
|
|
self.axis = [1]
|
|
self.axis_type = "int32"
|
|
|
|
def config_dtype(self):
|
|
self.x_type = "float64"
|
|
|
|
|
|
class TestGatherOp1FP16(TestGatherOp1):
|
|
def config_dtype(self):
|
|
self.x_type = "float16"
|
|
|
|
|
|
class TestGatherOp1Complex64(TestGatherOp1):
|
|
def config_dtype(self):
|
|
self.x_type = "complex64"
|
|
|
|
def test_check_grad(self):
|
|
self.check_grad(['X'], 'Out')
|
|
|
|
|
|
class TestGatherOp1Complex128(TestGatherOp1):
|
|
def config_dtype(self):
|
|
self.x_type = "complex128"
|
|
|
|
def test_check_grad(self):
|
|
self.check_grad(['X'], 'Out')
|
|
|
|
|
|
class TestGatherOp2(TestGatherOp1):
|
|
def config(self):
|
|
"""
|
|
For multi-dimension input
|
|
"""
|
|
self.x_shape = (10, 88, 10)
|
|
self.config_dtype()
|
|
self.index = [1, 3, 5]
|
|
self.index_type = "int64"
|
|
self.axis = [0]
|
|
self.axis_type = "int32"
|
|
|
|
def config_dtype(self):
|
|
self.x_type = "float64"
|
|
|
|
|
|
class TestGatherOp2FP16(TestGatherOp2):
|
|
def config_dtype(self):
|
|
self.x_type = "float16"
|
|
|
|
|
|
class TestGatherOp2Complex64(TestGatherOp2):
|
|
def config_dtype(self):
|
|
self.x_type = "complex64"
|
|
|
|
def test_check_grad(self):
|
|
self.check_grad(['X'], 'Out')
|
|
|
|
|
|
class TestGatherOp2Complex128(TestGatherOp2):
|
|
def config_dtype(self):
|
|
self.x_type = "complex128"
|
|
|
|
def test_check_grad(self):
|
|
self.check_grad(['X'], 'Out')
|
|
|
|
|
|
class TestGatherOp3(TestGatherOp1):
|
|
def config(self):
|
|
"""
|
|
For multi-dimension input
|
|
"""
|
|
self.x_shape = (10, 88, 10)
|
|
self.config_dtype()
|
|
self.index = [1, 3, 5]
|
|
self.index_type = "int64"
|
|
self.axis = [2]
|
|
self.axis_type = "int32"
|
|
|
|
def config_dtype(self):
|
|
self.x_type = "float64"
|
|
|
|
|
|
class TestGatherOp3FP16(TestGatherOp3):
|
|
def config_dtype(self):
|
|
self.x_type = "float16"
|
|
|
|
|
|
class TestGatherOp3Complex64(TestGatherOp3):
|
|
def config_dtype(self):
|
|
self.x_type = "complex64"
|
|
|
|
def test_check_grad(self):
|
|
self.check_grad(['X'], 'Out')
|
|
|
|
|
|
class TestGatherOp3Complex128(TestGatherOp3):
|
|
def config_dtype(self):
|
|
self.x_type = "complex128"
|
|
|
|
def test_check_grad(self):
|
|
self.check_grad(['X'], 'Out')
|
|
|
|
|
|
class TestGatherOp4(TestGatherOp1):
|
|
def config(self):
|
|
"""
|
|
For multi-dimension input
|
|
"""
|
|
self.x_shape = (3, 100, 10)
|
|
self.config_dtype()
|
|
self.index = [1, 1, 1, 1, 1, 1, 1, 1, 1, 1]
|
|
self.index_type = "int64"
|
|
self.axis = [0]
|
|
self.axis_type = "int32"
|
|
self.attrs = {'overwrite': False}
|
|
|
|
def config_dtype(self):
|
|
self.x_type = "float64"
|
|
|
|
|
|
class TestGatherOp4FP16(TestGatherOp4):
|
|
def config_dtype(self):
|
|
self.x_type = "float16"
|
|
|
|
|
|
class TestGatherOp4Complex64(TestGatherOp4):
|
|
def config_dtype(self):
|
|
self.x_type = "complex64"
|
|
|
|
def test_check_grad(self):
|
|
self.check_grad(['X'], 'Out')
|
|
|
|
|
|
class TestGatherOp4Complex128(TestGatherOp4):
|
|
def config_dtype(self):
|
|
self.x_type = "complex128"
|
|
|
|
def test_check_grad(self):
|
|
self.check_grad(['X'], 'Out')
|
|
|
|
|
|
class TestGatherOp5(TestGatherOp):
|
|
def config(self):
|
|
"""
|
|
Test for negative axis
|
|
"""
|
|
self.x_shape = (3, 100, 10)
|
|
self.config_dtype()
|
|
self.index = [1, 1, 1, 1, 1, 1, 1, 1, 1, 1]
|
|
self.index_type = "int64"
|
|
self.axis = [-1]
|
|
self.axis_type = "int32"
|
|
self.attrs = {'overwrite': False}
|
|
|
|
def config_dtype(self):
|
|
self.x_type = "float64"
|
|
|
|
def test_check_grad(self):
|
|
self.check_grad(
|
|
['X'],
|
|
'Out',
|
|
check_pir=True,
|
|
check_prim_pir=True,
|
|
)
|
|
|
|
|
|
class API_TestGather(unittest.TestCase):
|
|
def test_out1(self):
|
|
with base.program_guard(base.Program(), base.Program()):
|
|
data1 = paddle.static.data('data1', shape=[-1, 2], dtype='float64')
|
|
index = paddle.static.data('index', shape=[-1, 1], dtype='int64')
|
|
out = paddle.gather(data1, index)
|
|
place = base.CPUPlace()
|
|
exe = base.Executor(place)
|
|
input = np.array([[1, 2], [3, 4], [5, 6]]).astype('float64')
|
|
index_1 = np.array([1, 2]).astype('int64')
|
|
(result,) = exe.run(
|
|
feed={"data1": input, "index": index_1}, fetch_list=[out]
|
|
)
|
|
expected_output = np.array([[3, 4], [5, 6]])
|
|
np.testing.assert_allclose(result, expected_output, rtol=1e-05)
|
|
|
|
def test_out2(self):
|
|
with paddle.static.program_guard(
|
|
paddle.static.Program(), paddle.static.Program()
|
|
):
|
|
x = paddle.static.data('x', shape=[-1, 2], dtype='float64')
|
|
index = paddle.static.data('index', shape=[-1, 1], dtype='int32')
|
|
axis = paddle.static.data('axis', shape=[1], dtype='int32')
|
|
out = paddle.gather(x, index, axis)
|
|
place = paddle.CPUPlace()
|
|
exe = paddle.static.Executor(place)
|
|
x_np = np.array([[1, 2], [3, 4], [5, 6]]).astype('float64')
|
|
index_np = np.array([1, 1]).astype('int32')
|
|
axis_np = np.array([1]).astype('int32')
|
|
(result,) = exe.run(
|
|
feed={"x": x_np, "index": index_np, 'axis': axis_np},
|
|
fetch_list=[out],
|
|
)
|
|
expected_output = gather_numpy(x_np, index_np, axis_np[0])
|
|
np.testing.assert_allclose(result, expected_output, rtol=1e-05)
|
|
|
|
|
|
@unittest.skipIf(
|
|
not (core.is_compiled_with_cuda() or is_custom_device()),
|
|
"only support compiled with CUDA.",
|
|
)
|
|
class TestGatherGPUCPUConsistency(unittest.TestCase):
|
|
def test_gpu_cpu_consistency(self):
|
|
paddle.disable_static()
|
|
np.random.seed(42)
|
|
x = np.random.rand(1000, 128).astype("float32")
|
|
index = np.random.randint(0, 1000, size=(100,))
|
|
cpu_out = paddle.gather(
|
|
paddle.to_tensor(x, place=paddle.CPUPlace()),
|
|
paddle.to_tensor(index),
|
|
)
|
|
gpu_out = paddle.gather(
|
|
paddle.to_tensor(x, place=paddle.CUDAPlace(0)),
|
|
paddle.to_tensor(index),
|
|
)
|
|
np.testing.assert_allclose(cpu_out.numpy(), gpu_out.numpy(), rtol=1e-6)
|
|
paddle.enable_static()
|
|
|
|
|
|
class API_TestDygraphGather(unittest.TestCase):
|
|
def test_out1(self):
|
|
paddle.disable_static()
|
|
input_1 = np.array([[1, 2], [3, 4], [5, 6]])
|
|
index_1 = np.array([1, 2])
|
|
input = paddle.to_tensor(input_1)
|
|
index = paddle.to_tensor(index_1)
|
|
output = paddle.gather(input, index)
|
|
output_np = output.numpy()
|
|
expected_output = np.array([[3, 4], [5, 6]])
|
|
np.testing.assert_allclose(output_np, expected_output, rtol=1e-05)
|
|
paddle.enable_static()
|
|
|
|
def test_out12(self):
|
|
paddle.disable_static()
|
|
input_1 = np.array([[1, 2], [3, 4], [5, 6]])
|
|
index_1 = np.array([1, 2])
|
|
x = paddle.to_tensor(input_1)
|
|
index = paddle.to_tensor(index_1)
|
|
output = paddle.gather(x, index, axis=0)
|
|
output_np = output.numpy()
|
|
expected_output = gather_numpy(input_1, index_1, axis=0)
|
|
np.testing.assert_allclose(output_np, expected_output, rtol=1e-05)
|
|
paddle.enable_static()
|
|
|
|
def test_zero_index(self):
|
|
paddle.disable_static()
|
|
x = paddle.to_tensor([[1, 2], [3, 4]])
|
|
index = paddle.to_tensor(np.array([]).astype('int64'))
|
|
for axis in range(len(x.shape)):
|
|
out = paddle.gather(x, index, axis)
|
|
expected_shape = list(x.shape)
|
|
expected_shape[axis] = 0
|
|
self.assertEqual(list(out.shape), expected_shape)
|
|
paddle.enable_static()
|
|
|
|
def test_large_data(self):
|
|
if not (paddle.is_compiled_with_cuda() or is_custom_device()):
|
|
return
|
|
|
|
x = np.random.rand(226862, 256).astype("float32")
|
|
index = np.random.randint(-226862, 22682, size=(8859027))
|
|
|
|
def test_dygraph():
|
|
with base.dygraph.guard():
|
|
gpu_out = paddle.gather(
|
|
paddle.to_tensor(x), paddle.to_tensor(index)
|
|
)
|
|
return gpu_out.numpy()
|
|
|
|
@switch_to_static_graph
|
|
def test_static_graph():
|
|
with paddle.static.program_guard(
|
|
paddle.static.Program(), paddle.static.Program()
|
|
):
|
|
x_t = paddle.static.data(name="x", dtype=x.dtype, shape=x.shape)
|
|
index_t = paddle.static.data(
|
|
name="index", dtype=index.dtype, shape=index.shape
|
|
)
|
|
out_t = paddle.gather(x_t, index_t)
|
|
feed = {x_t.name: x, index_t.name: index}
|
|
fetch = [out_t]
|
|
|
|
gpu_exe = paddle.static.Executor(get_device_place())
|
|
gpu_value = gpu_exe.run(feed=feed, fetch_list=fetch)[0]
|
|
return gpu_value
|
|
|
|
np.testing.assert_array_equal(test_dygraph(), test_static_graph())
|
|
|
|
|
|
class TestGathertError(unittest.TestCase):
|
|
def test_error1(self):
|
|
with paddle.static.program_guard(
|
|
paddle.static.Program(), paddle.static.Program()
|
|
):
|
|
shape = [8, 9, 6]
|
|
x = paddle.static.data(shape=shape, dtype='int8', name='x')
|
|
axis = paddle.static.data(shape=[1], dtype='float32', name='axis')
|
|
index = paddle.static.data(shape=shape, dtype='int32', name='index')
|
|
index_float = paddle.static.data(
|
|
shape=shape, dtype='float32', name='index_float'
|
|
)
|
|
|
|
def test_x_type():
|
|
paddle.gather(x, index)
|
|
|
|
self.assertRaises((TypeError, ValueError), test_x_type)
|
|
|
|
def test_index_type():
|
|
paddle.gather(x, index_float)
|
|
|
|
self.assertRaises((TypeError, ValueError), test_index_type)
|
|
|
|
def test_axis_dtype():
|
|
paddle.gather(x, index, axis=1.11)
|
|
|
|
self.assertRaises((TypeError, ValueError), test_axis_dtype)
|
|
|
|
def test_axis_dtype1():
|
|
paddle.gather(x, index, axis=axis)
|
|
|
|
self.assertRaises((TypeError, ValueError), test_axis_dtype1)
|
|
|
|
def test_error2(self):
|
|
with paddle.static.program_guard(
|
|
paddle.static.Program(), paddle.static.Program()
|
|
):
|
|
shape = [8, 9, 6]
|
|
x = paddle.static.data(shape=shape, dtype='int8', name='x')
|
|
index = paddle.static.data(shape=shape, dtype='int32', name='mask')
|
|
index_float = paddle.static.data(
|
|
shape=shape, dtype='float32', name='index_float'
|
|
)
|
|
|
|
def test_x_type():
|
|
paddle.gather(x, index)
|
|
|
|
self.assertRaises((TypeError, ValueError), test_x_type)
|
|
|
|
def test_index_type():
|
|
paddle.gather(x, index_float)
|
|
|
|
self.assertRaises((TypeError, ValueError), test_index_type)
|
|
|
|
def test_error3(self):
|
|
with paddle.static.program_guard(
|
|
paddle.static.Program(), paddle.static.Program()
|
|
):
|
|
shape = [8, 9, 6]
|
|
x = paddle.static.data(shape=shape, dtype='int32', name='x')
|
|
index = paddle.static.data(shape=shape, dtype='int32', name='index')
|
|
|
|
def test_axis_minsize():
|
|
paddle.gather(x, index, axis=-1)
|
|
|
|
self.assertRaises(ValueError, test_axis_minsize)
|
|
|
|
def test_axis_maxsize():
|
|
paddle.gather(x, index, axis=512)
|
|
|
|
self.assertRaises(ValueError, test_axis_maxsize)
|
|
|
|
|
|
class TestCheckOutType(unittest.TestCase):
|
|
def test_out_type(self):
|
|
data = paddle.static.data(shape=[16, 10], dtype='int64', name='x')
|
|
index = paddle.static.data(shape=[4], dtype='int64', name='index')
|
|
out = paddle.gather(data, index)
|
|
self.assertTrue(
|
|
out.dtype == paddle.int64 or out.dtype == core.DataType.INT64
|
|
)
|
|
|
|
def test_pir_out_type(self):
|
|
with paddle.pir_utils.IrGuard():
|
|
data = paddle.static.data(shape=[16, 10], dtype='int64', name='x')
|
|
index = paddle.static.data(shape=[4], dtype='int64', name='index')
|
|
out = paddle.gather(data, index)
|
|
self.assertTrue(out.dtype == core.DataType.INT64)
|
|
|
|
|
|
class TestGatherBackward(unittest.TestCase):
|
|
def setUp(self):
|
|
self.shape = [10, 20]
|
|
self.dtype = 'float32'
|
|
self.index = (1, 3, 5)
|
|
self.index_dtype = 'int64'
|
|
self.places = get_places()
|
|
|
|
def test_gather_backward(self):
|
|
if len(self.places) != 2:
|
|
return
|
|
res_list = []
|
|
x_np = np.random.random(self.shape).astype(self.dtype)
|
|
index_np = np.array(self.index, dtype=self.index_dtype)
|
|
grad_out_np = np.random.random(self.shape).astype(self.dtype)
|
|
for place in self.places:
|
|
with base.dygraph.guard(place):
|
|
x = paddle.to_tensor(x_np, dtype=self.dtype)
|
|
x.stop_gradient = False
|
|
index = paddle.to_tensor(index_np, dtype=self.index_dtype)
|
|
out = paddle.gather(x, index, -1)
|
|
grad_out = paddle.to_tensor(grad_out_np, dtype=self.dtype)
|
|
(re,) = paddle.grad(
|
|
outputs=out,
|
|
inputs=x,
|
|
grad_outputs=grad_out,
|
|
)
|
|
res_list.append(re.numpy())
|
|
np.testing.assert_allclose(res_list[0], res_list[1])
|
|
|
|
|
|
class TestGatherOp_ZeroSize(OpTest):
|
|
def setUp(self):
|
|
self.op_type = "gather"
|
|
self.python_api = paddle.gather
|
|
self.public_python_api = paddle.gather
|
|
self.config()
|
|
self.init_inputs_and_outputs()
|
|
|
|
def test_check_output(self):
|
|
self.check_output(check_pir=True)
|
|
|
|
def test_check_grad(self):
|
|
self.check_grad(['X'], 'Out', check_pir=True)
|
|
|
|
def config(self):
|
|
self.x_shape = (3, 0, 4)
|
|
self.config_dtype()
|
|
self.index = [2]
|
|
self.index_type = "int32"
|
|
|
|
def config_dtype(self):
|
|
self.x_type = "float64"
|
|
|
|
def init_inputs_and_outputs(self):
|
|
xnp = np.random.random(self.x_shape).astype(self.x_type)
|
|
self.inputs = {
|
|
'X': xnp,
|
|
'Index': np.array(self.index).astype(self.index_type),
|
|
}
|
|
self.outputs = {'Out': self.inputs["X"][self.inputs["Index"]]}
|
|
|
|
|
|
class TestGatherOp_ZeroSize2(TestGatherOp_ZeroSize):
|
|
def config(self):
|
|
self.x_shape = (10, 20)
|
|
self.config_dtype()
|
|
self.index = [2, 0]
|
|
self.index_type = "int32"
|
|
|
|
|
|
if __name__ == "__main__":
|
|
paddle.enable_static()
|
|
unittest.main()
|