fix take_along_axis when -1 in arr.shape (#72436)

HydrogenSulfate · web-flow · commit 441816abec7d · 2025-04-25T17:47:20.000+08:00
diff --git a/python/paddle/tensor/manipulation.py b/python/paddle/tensor/manipulation.py
@@ -6558,7 +6558,11 @@ def take_along_axis(
         arr = paddle.broadcast_to(arr, broadcast_shape)
     else:
         for i in range(len(arr.shape)):
-            if i != axis and arr.shape[i] < indices.shape[i]:
+            if (
+                i != axis
+                and arr.shape[i] != -1
+                and arr.shape[i] < indices.shape[i]
+            ):
                 raise RuntimeError(
                     f"Size does not match at dimension {i} expected index {indices.shape} to be smaller than self {arr.shape} apart from dimension {axis}"
                 )
diff --git a/test/legacy_test/test_take_along_axis_op.py b/test/legacy_test/test_take_along_axis_op.py
@@ -18,6 +18,7 @@
 
 import numpy as np
 from op_test import OpTest, convert_float_to_uint16
+from utils import dygraph_guard
 
 import paddle
 from paddle.framework import core
@@ -378,6 +379,28 @@ def test_error(self):
             res = paddle.take_along_axis(tensorx, indices, 0, False)
 
 
+class TestTakeAlongAxisAPICase4(unittest.TestCase):
+    def test_static_shape_take_along_axis(self):
+        with dygraph_guard():
+
+            x = paddle.randn([4, 2])
+            ind = paddle.to_tensor([[0, 1]])
+
+            static_f = paddle.jit.to_static(
+                paddle.take_along_axis,
+                input_spec=[
+                    paddle.static.InputSpec(
+                        shape=[-1, -1], dtype="float32", name="arr"
+                    ),
+                    paddle.static.InputSpec(
+                        shape=[-1, 2], dtype="int64", name="indices"
+                    ),
+                ],
+                full_graph=True,
+            )
+            _ = static_f(x, ind, axis=0, broadcast=False)
+
+
 if __name__ == "__main__":
     paddle.enable_static()
     unittest.main()
