update nn.module test

Differential Revision: D58907317

Author: huiyujie0105

backends/vulkan/test/test_vulkan_delegate.py

@@ -26,6 +26,7 @@
 ctypes.CDLL("libvulkan.so.1")
 
 
+from executorch.exir.program._program import _to_edge_transform_and_lower
 from executorch.extension.pybindings.portable_lib import (  # @manual
     _load_for_executorch_from_buffer,
 )
@@ -100,6 +101,7 @@ def lower_module_and_test_output(
         test_inputs=None,
         memory_layouts=None,
         first_output_only=False,
+        not_decompose=False,
     ):
         """
         Helper testing function that takes a torch.nn.Module and lowers it to Vulkan with
@@ -119,16 +121,25 @@ def run_test(memory_layout):
             program: ExportedProgram = export(
                 model, sample_inputs, dynamic_shapes=dynamic_shapes
             )
-            edge_program: EdgeProgramManager = to_edge(
-                program, compile_config=self._edge_compile_config
-            )
-
-            edge_program = edge_program.transform(
-                [I64toI32(self._edge_compile_config._skip_dim_order), MeanToSumDiv()]
-            )
+            if not_decompose is False:
+                edge_program: EdgeProgramManager = to_edge(
+                    program, compile_config=self._edge_compile_config
+                )
 
-            edge_program = edge_program.to_backend(VulkanPartitioner(compile_options))
+                edge_program = edge_program.transform(
+                    [
+                        I64toI32(self._edge_compile_config._skip_dim_order),
+                        MeanToSumDiv(),
+                    ]
+                )
 
+                edge_program = edge_program.to_backend(
+                    VulkanPartitioner(compile_options)
+                )
+            else:
+                edge_program = _to_edge_transform_and_lower(
+                    program, partitioner=[VulkanPartitioner(compile_options)]
+                )
             executorch_program = edge_program.to_executorch()
 
             self.assertEqual(