def do_eliminate_transpose(test_case, with_cuda):
x = flow.randn(2, 3, 4, 5)
if with_cuda:
x = x.cuda()
eager_res = flow.permute(flow.permute(x, (0, 2, 3, 1)), (0, 3, 1, 2))
class GraphToRun(flow.nn.Graph):
def __init__(self):
super().__init__()
def build(self, x):
return flow.permute(flow.permute(x, (0, 2, 3, 1)), (0, 3, 1, 2))
graph_to_run = GraphToRun()
lazy_res = graph_to_run(x)
test_case.assertTrue(
np.allclose(eager_res.numpy(), lazy_res.numpy(), rtol=1e-5, atol=1e-5))
def _test_permute_impl(test_case, device):
input = flow.tensor(
np.random.randn(2, 6, 5, 3),
dtype=flow.float32,
device=flow.device(device),
requires_grad=True,
)
of_out1 = flow.permute(input, (1, 0, 2, 3))
np_out = input.numpy().transpose((1, 0, 2, 3))
test_case.assertTrue(np.array_equal(of_out1.numpy().flatten(), np_out.flatten()))
of_out = of_out1.sum()
of_out.backward()
np_grad = np.ones((2, 6, 5, 3))
test_case.assertTrue(np.allclose(input.grad.numpy(), np_grad, 0.0001, 0.0001))
def build(self, x):
return flow.permute(flow.permute(x, (0, 2, 3, 1)), (0, 3, 1, 2))