【PIR API adaptor No.250】 python/paddle/geometric/sampling/neighbors.py (#58783)

Liyulingyue · web-flow · commit 69fcbf8d842f · 2023-11-10T17:07:30.000+08:00
diff --git a/python/paddle/geometric/sampling/neighbors.py b/python/paddle/geometric/sampling/neighbors.py
@@ -15,7 +15,7 @@
 from paddle import _C_ops, _legacy_C_ops
 from paddle.base.data_feeder import check_variable_and_dtype
 from paddle.base.layer_helper import LayerHelper
-from paddle.framework import in_dynamic_mode
+from paddle.framework import in_dynamic_mode, in_dynamic_or_pir_mode
 
 __all__ = []
 
@@ -251,7 +251,7 @@ def weighted_sample_neighbors(
                 "`eids` should not be None if `return_eids` is True."
             )
 
-    if in_dynamic_mode():
+    if in_dynamic_or_pir_mode():
         (
             out_neighbors,
             out_count,
diff --git a/test/legacy_test/test_weighted_sample_neighbors.py b/test/legacy_test/test_weighted_sample_neighbors.py
@@ -17,6 +17,7 @@
 import numpy as np
 
 import paddle
+from paddle.pir_utils import test_with_pir_api
 
 
 class TestWeightedSampleNeighbors(unittest.TestCase):
@@ -80,6 +81,7 @@ def test_sample_result(self):
             )
             self.assertTrue(np.sum(in_neighbors) == in_neighbors.shape[0])
 
+    @test_with_pir_api
     def test_sample_result_static(self):
         paddle.enable_static()
         with paddle.static.program_guard(paddle.static.Program()):

Original file line number	Diff line number	Diff line change
`@@ -15,7 +15,7 @@`
`15`	`15`	`from paddle import _C_ops, _legacy_C_ops`
`16`	`16`	`from paddle.base.data_feeder import check_variable_and_dtype`
`17`	`17`	`from paddle.base.layer_helper import LayerHelper`
`18`		`-from paddle.framework import in_dynamic_mode`
	`18`	`+from paddle.framework import in_dynamic_mode, in_dynamic_or_pir_mode`
`19`	`19`
`20`	`20`	`__all__ = []`
`21`	`21`
`@@ -251,7 +251,7 @@ def weighted_sample_neighbors(`
`251`	`251`	"`eids` should not be None if `return_eids` is True."
`252`	`252`	`)`
`253`	`253`
`254`		`- if in_dynamic_mode():`
	`254`	`+ if in_dynamic_or_pir_mode():`
`255`	`255`	`(`
`256`	`256`	`out_neighbors,`
`257`	`257`	`out_count,`