【PIR API adaptor No.177】Migrate paddle.geometric.reindex_graph into pir (#60323)

ooooo-create · web-flow · commit fdf9d71ed6b9 · 2023-12-26T16:00:25.000+08:00
diff --git a/python/paddle/geometric/reindex.py b/python/paddle/geometric/reindex.py
@@ -17,7 +17,7 @@
 from paddle.base.data_feeder import check_variable_and_dtype
 from paddle.base.framework import Variable
 from paddle.base.layer_helper import LayerHelper
-from paddle.framework import in_dynamic_mode, in_dynamic_or_pir_mode
+from paddle.framework import in_dynamic_or_pir_mode
 
 __all__ = []
 
@@ -90,7 +90,7 @@ def reindex_graph(
         True if value_buffer is not None and index_buffer is not None else False
     )
 
-    if in_dynamic_mode():
+    if in_dynamic_or_pir_mode():
         reindex_src, reindex_dst, out_nodes = _C_ops.reindex_graph(
             x,
             neighbors,
diff --git a/test/legacy_test/test_graph_reindex.py b/test/legacy_test/test_graph_reindex.py
@@ -370,6 +370,7 @@ def test_heter_reindex_result_v3(self):
         np.testing.assert_allclose(reindex_dst, reindex_dst_, rtol=1e-05)
         np.testing.assert_allclose(out_nodes, out_nodes_, rtol=1e-05)
 
+    @test_with_pir_api
     def test_reindex_result_static(self):
         paddle.enable_static()
         with paddle.static.program_guard(paddle.static.Program()):
