From 81c4caa7ab12a91ae28eb235ba1112c36f989fb9 Mon Sep 17 00:00:00 2001 From: Tugsbayasgalan Manlaibaatar Date: Tue, 5 Nov 2024 10:13:08 -0700 Subject: [PATCH] Make XLA tests export new IR compatible (#8354) --- test/stablehlo/test_export_fx_passes.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/test/stablehlo/test_export_fx_passes.py b/test/stablehlo/test_export_fx_passes.py index 82650997316..580b756e177 100644 --- a/test/stablehlo/test_export_fx_passes.py +++ b/test/stablehlo/test_export_fx_passes.py @@ -212,7 +212,8 @@ def forward(self, x, dim, weight, bias, eps): after_decomp_out = native_layer_norm_impl(*args) self.assertTrue( torch.allclose(before_decomp_out, after_decomp_out, atol=1e-6)) - ep = export(m, args, dynamic_shapes=dynamic_shapes) + ep_training = export(m, args, dynamic_shapes=dynamic_shapes) + ep = ep_training.run_decompositions({}) decompose_dynamic_native_layer_norm(ep.graph_module) ep.graph_module.recompile() self.assertFalse('aten.native_layer_norm' in ep.graph_module.code)