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

Author: Liyulingyue

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,

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()):