Fix crash of transpose

cybersupersoap · cybersupersoap · commit e9009cefdb30 · 2025-02-19T21:37:20.000+08:00
diff --git a/tensorflow/core/kernels/transpose_op.cc b/tensorflow/core/kernels/transpose_op.cc
@@ -150,7 +150,10 @@ void TransposeOp::Compute(OpKernelContext* ctx) {
   bool is_identity = true;
   for (int i = 0; i < dims; ++i) {
     int32_t d = permutation[i];
-    if (d < 0) d += dims;
+    if (d < 0) {
+      d += dims;
+      permutation[i] = d;
+    }
     OP_REQUIRES(
         ctx, 0 <= d && d < dims,
         errors::InvalidArgument(d, " is out of range [0 .. ", dims, ")"));
diff --git a/tensorflow/python/kernel_tests/math_ops/transpose_op_test.py b/tensorflow/python/kernel_tests/math_ops/transpose_op_test.py
@@ -541,6 +541,13 @@ def testError(self):
     self._testError(
         np.arange(0., 30).reshape([2, 3, 5]), [0, 1, 1], "2 is missing")
 
+  def testNegativePerm(self):
+    self.assertEqual([15, 100, 37],
+                    array_ops.transpose(
+                           constant_op.constant(
+                               1, dtype=dtypes.int32, shape=[100, 37, 15]),
+                           [-1, -3, -2]).get_shape().dims)
+
 
 if __name__ == "__main__":
   test.main()
