Skip to content

Commit

Permalink
[PIR] No.28 Migrate paddle.gather_nd into pir (PaddlePaddle#57562)
Browse files Browse the repository at this point in the history 
* support gather_nd in pir

* update test

* add false

* fix bug
  • Loading branch information
@zrr1999 @jiahy0825
zrr1999 authored and jiahy0825 committed Oct 16, 2023
1 parent 4d1e9d7 commit ad5860d
Show file tree
Hide file tree
Showing 2 changed files with 58 additions and 29 deletions.
2 changes: 1 addition & 1 deletion python/paddle/tensor/manipulation.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3798,7 +3798,7 @@ def gather_nd(x, index, name=None):
output = paddle.gather_nd(x, index) #[[3, 4]]
"""
if in_dynamic_mode():
if in_dynamic_or_pir_mode():
return _C_ops.gather_nd(x, index)
else:
check_variable_and_dtype(
Expand Down
85 changes: 57 additions & 28 deletions test/legacy_test/test_gather_nd_op.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -53,10 +53,10 @@ def config_dtype(self):
self.dtype = np.float64

def test_check_output(self):
self.check_output()
self.check_output(check_new_ir=True)

def test_check_grad(self):
self.check_grad(['X'], 'Out', check_prim=True)
self.check_grad(['X'], 'Out', check_prim=True, check_new_ir=False)


class TestGatherNdOpWithEmptyIndexFP16(TestGatherNdOpWithEmptyIndex):
Expand All @@ -75,11 +75,13 @@ def config_dtype(self):

def test_check_output(self):
place = core.CUDAPlace(0)
self.check_output_with_place(place)
self.check_output_with_place(place, check_new_ir=True)

def test_check_grad(self):
place = core.CUDAPlace(0)
self.check_grad_with_place(place, ['X'], 'Out', check_prim=True)
self.check_grad_with_place(
place, ['X'], 'Out', check_prim=True, check_new_ir=False
)


class TestGatherNdOpWithIndex1(OpTest):
Expand Down Expand Up @@ -112,10 +114,10 @@ def config_dtype(self):
self.dtype = np.float64

def test_check_output(self):
self.check_output()
self.check_output(check_new_ir=True)

def test_check_grad(self):
self.check_grad(['X'], 'Out', check_prim=True)
self.check_grad(['X'], 'Out', check_prim=True, check_new_ir=False)


class TestGatherNdOpWithIndex1_ZeroDim(TestGatherNdOpWithIndex1):
Expand Down Expand Up @@ -161,11 +163,13 @@ def config_dtype(self):

def test_check_output(self):
place = core.CUDAPlace(0)
self.check_output_with_place(place)
self.check_output_with_place(place, check_new_ir=True)

def test_check_grad(self):
place = core.CUDAPlace(0)
self.check_grad_with_place(place, ['X'], 'Out', check_prim=True)
self.check_grad_with_place(
place, ['X'], 'Out', check_prim=True, check_new_ir=False
)


class TestGatherNdOpWithLowIndex(OpTest):
Expand Down Expand Up @@ -198,10 +202,10 @@ def config_dtype(self):
self.dtype = np.float64

def test_check_output(self):
self.check_output()
self.check_output(check_new_ir=True)

def test_check_grad(self):
self.check_grad(['X'], 'Out', check_prim=True)
self.check_grad(['X'], 'Out', check_prim=True, check_new_ir=False)


class TestGatherNdOpWithLowIndexFP16(TestGatherNdOpWithLowIndex):
Expand All @@ -220,12 +224,17 @@ def config_dtype(self):

def test_check_output(self):
place = core.CUDAPlace(0)
self.check_output_with_place(place)
self.check_output_with_place(place, check_new_ir=True)

def test_check_grad(self):
place = core.CUDAPlace(0)
self.check_grad_with_place(
place, ['X'], 'Out', check_prim=True, numeric_grad_delta=0.5
place,
['X'],
'Out',
check_prim=True,
check_new_ir=False,
numeric_grad_delta=0.5,
)


Expand Down Expand Up @@ -264,10 +273,16 @@ def config_dtype(self):
self.dtype = np.float64

def test_check_output(self):
self.check_output()
self.check_output(check_new_ir=True)

def test_check_grad(self):
self.check_grad(['X'], 'Out', check_prim=True, numeric_grad_delta=0.05)
self.check_grad(
['X'],
'Out',
check_prim=True,
check_new_ir=False,
numeric_grad_delta=0.05,
)


class TestGatherNdOpIndex1FP16(TestGatherNdOpIndex1):
Expand All @@ -286,12 +301,17 @@ def config_dtype(self):

def test_check_output(self):
place = core.CUDAPlace(0)
self.check_output_with_place(place)
self.check_output_with_place(place, check_new_ir=True)

def test_check_grad(self):
place = core.CUDAPlace(0)
self.check_grad_with_place(
place, ['X'], 'Out', check_prim=True, numeric_grad_delta=0.5
place,
['X'],
'Out',
check_prim=True,
check_new_ir=False,
numeric_grad_delta=0.5,
)


Expand Down Expand Up @@ -322,10 +342,10 @@ def config_dtype(self):
self.dtype = np.float64

def test_check_output(self):
self.check_output()
self.check_output(check_new_ir=True)

def test_check_grad(self):
self.check_grad(['X'], 'Out', check_prim=True)
self.check_grad(['X'], 'Out', check_prim=True, check_new_ir=False)


class TestGatherNdOpWithSameIndexAsXFP16(TestGatherNdOpWithSameIndexAsX):
Expand All @@ -344,12 +364,17 @@ def config_dtype(self):

def test_check_output(self):
place = core.CUDAPlace(0)
self.check_output_with_place(place)
self.check_output_with_place(place, check_new_ir=True)

def test_check_grad(self):
place = core.CUDAPlace(0)
self.check_grad_with_place(
place, ['X'], 'Out', check_prim=True, numeric_grad_delta=0.5
place,
['X'],
'Out',
check_prim=True,
check_new_ir=False,
numeric_grad_delta=0.5,
)


Expand Down Expand Up @@ -382,10 +407,10 @@ def config_dtype(self):
self.dtype = np.float64

def test_check_output(self):
self.check_output()
self.check_output(check_new_ir=True)

def test_check_grad(self):
self.check_grad(['X'], 'Out', check_prim=True)
self.check_grad(['X'], 'Out', check_prim=True, check_new_ir=False)


class TestGatherNdOpWithHighRankSameFP16(TestGatherNdOpWithHighRankSame):
Expand All @@ -404,11 +429,13 @@ def config_dtype(self):

def test_check_output(self):
place = core.CUDAPlace(0)
self.check_output_with_place(place)
self.check_output_with_place(place, check_new_ir=True)

def test_check_grad(self):
place = core.CUDAPlace(0)
self.check_grad_with_place(place, ['X'], 'Out', check_prim=True)
self.check_grad_with_place(
place, ['X'], 'Out', check_prim=True, check_new_ir=False
)


class TestGatherNdOpWithHighRankDiff(OpTest):
Expand Down Expand Up @@ -441,10 +468,10 @@ def config_dtype(self):
self.dtype = np.float64

def test_check_output(self):
self.check_output()
self.check_output(check_new_ir=True)

def test_check_grad(self):
self.check_grad(['X'], 'Out', check_prim=True)
self.check_grad(['X'], 'Out', check_prim=True, check_new_ir=False)


class TestGatherNdOpWithHighRankDiffFP16(TestGatherNdOpWithHighRankDiff):
Expand All @@ -463,11 +490,13 @@ def config_dtype(self):

def test_check_output(self):
place = core.CUDAPlace(0)
self.check_output_with_place(place)
self.check_output_with_place(place, check_new_ir=True)

def test_check_grad(self):
place = core.CUDAPlace(0)
self.check_grad_with_place(place, ['X'], 'Out', check_prim=True)
self.check_grad_with_place(
place, ['X'], 'Out', check_prim=True, check_new_ir=False
)


# Test Python API
Expand Down

0 comments on commit ad5860d

Please sign in to comment.