remove program_guard

Author: zhiqiu

python/paddle/distributed/auto_parallel/static/pir_pass.py

@@ -165,22 +165,24 @@ def apply_reshard_pass(program):
 # training, the transpose is equal to reshape.
 # So, this pass is to haddle the specific case.
 def eliminate_transpose_by_reshape(program):
-    with paddle.static.program_guard(program):
-        for op in program.global_block().ops:
-            if op.name() == 'pd_op.transpose' or op.name() == 'pd_op.transpose':
-                var = op.operand(0).source()
-                rank = len(var.shape)
-                perm = op.attrs()['perm']
-                perm = [p + rank if p < 0 else p for p in perm]
-                # only support transpose dim 0 and dim 1
-                expected_perm = [1, 0] + [i + 2 for i in range(rank - 2)]
-                if perm == expected_perm and (
-                    var.shape[0] == 1 or var.shape[1] == 1
-                ):
-                    if var.shape == [1, 1024, 4096]:
-                        paddle.pir.set_insertion_point(op)
-                        transpose_var = op.result(0)
-                        reshape_var = paddle.reshape(var, transpose_var.shape)
-                        transpose_var.replace_all_uses_with(reshape_var)
-                        program.global_block().remove_op(op)
+    for op in program.global_block().ops:
+        if (
+            op.name() == 'pd_op.transpose'
+            or op.name() == 'pd_op.transpose_grad'
+        ):
+            var = op.operand(0).source()
+            rank = len(var.shape)
+            perm = op.attrs()['perm']
+            perm = [p + rank if p < 0 else p for p in perm]
+            # only support transpose dim 0 and dim 1
+            expected_perm = [1, 0] + [i + 2 for i in range(rank - 2)]
+            if perm == expected_perm and (
+                var.shape[0] == 1 or var.shape[1] == 1
+            ):
+                print('elinimate', op)
+                paddle.pir.set_insertion_point(op)
+                transpose_var = op.result(0)
+                reshape_var = paddle._C_ops.reshape(var, transpose_var.shape)
+                transpose_var.replace_all_uses_with(reshape_var)
+                program.global_block().remove_op(op)
     return program