Arm backend: Fix Gemma3n and CLIP export tests (pytorch#22830)

Index Gemma3n audio encoder outputs to handle both tuples and
ModelOutput. Run CLIP eagerly during test setup to install hidden-state
hooks before strict export encounters their installation lock.

cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils
@Sebastian-Larsson @robell @rascani

Signed-off-by: Sangwon Ha <sangwon.ha@arm.com>

Author: FabulousSuperDude

backends/arm/test/models/Gemma3n/test_gemma3nModel.py

@@ -548,8 +548,7 @@ def __init__(self, config) -> None:
         self.encoder = Gemma3nAudioEncoder(config)
 
     def forward(self, x: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
-        encodings, _ = self.encoder(x, mask)
-        return encodings
+        return self.encoder(x, mask)[0]
 
     @staticmethod
     def _prepare_inputs(

backends/arm/test/models/stable_diffusion_3_5_large/test_CLIPTextModelWithProjection_sd35_large.py

@@ -67,14 +67,19 @@ def create_dummy_inputs(
             ),
         )
 
-    def create_model(
+    def prepare_model_and_inputs(
         self,
         config,
-    ) -> SD3CLIPTextEncoderWrapper:
-        """Instantiate wrapped CLIPTextModelWithProjection for tests."""
-        return SD3CLIPTextEncoderWrapper(
+    ) -> tuple[SD3CLIPTextEncoderWrapper, input_t]:
+        """Prepare wrapped CLIPTextModelWithProjection for export."""
+        model = SD3CLIPTextEncoderWrapper(
             CLIPTextModelWithProjection(config).to(dtype=config.dtype)  # type: ignore[call-arg]
         ).eval()
+        inputs = self.create_dummy_inputs(config)
+        # Install Transformers' lazy hidden-state hooks before strict export.
+        with torch.no_grad():
+            model(*inputs)
+        return model, inputs
 
     @staticmethod
     def ops_after_partitioner_INT(config) -> dict[str, int]:
@@ -107,11 +112,12 @@ def test_clip_text_model_with_projection_tosa_FP(config_factory, atol):
     """Run the CLIPTextModelWithProjection TOSA FP test for a given config."""
     test_helper = TestCLIPTextModelWithProjection()
     config = config_factory()
+    model, inputs = test_helper.prepare_model_and_inputs(config)
 
     with torch.no_grad():
         pipeline = TosaPipelineFP[input_t](
-            test_helper.create_model(config),
-            test_helper.create_dummy_inputs(config),
+            model,
+            inputs,
             aten_op=[],
             exir_op=[],
             use_to_edge_transform_and_lower=True,
@@ -136,11 +142,12 @@ def test_clip_text_model_with_projection_tosa_INT(config_factory, atol):
     """Run the CLIPTextModelWithProjection TOSA INT test for a given config."""
     test_helper = TestCLIPTextModelWithProjection()
     config = config_factory()
+    model, inputs = test_helper.prepare_model_and_inputs(config)
 
     with torch.no_grad():
         pipeline = TosaPipelineINT[input_t](
-            test_helper.create_model(config),
-            test_helper.create_dummy_inputs(config),
+            model,
+            inputs,
             aten_op=[],
             exir_op=[],
             use_to_edge_transform_and_lower=True,
@@ -168,11 +175,12 @@ def test_clip_text_model_with_projection_vgf_no_quant(config_factory):
     """Run the CLIPTextModelWithProjection VGF no-quant test."""
     test_helper = TestCLIPTextModelWithProjection()
     config = config_factory()
+    model, inputs = test_helper.prepare_model_and_inputs(config)
 
     with torch.no_grad():
         pipeline = VgfPipeline[input_t](
-            test_helper.create_model(config),
-            test_helper.create_dummy_inputs(config),
+            model,
+            inputs,
             aten_op=[],
             exir_op=[],
             use_to_edge_transform_and_lower=True,
@@ -200,11 +208,12 @@ def test_clip_text_model_with_projection_vgf_quant(config_factory, atol):
     """Run the CLIPTextModelWithProjection VGF quant test."""
     test_helper = TestCLIPTextModelWithProjection()
     config = config_factory()
+    model, inputs = test_helper.prepare_model_and_inputs(config)
 
     with torch.no_grad():
         pipeline = VgfPipeline[input_t](
-            test_helper.create_model(config),
-            test_helper.create_dummy_inputs(config),
+            model,
+            inputs,
             aten_op=[],
             exir_op=[],
             use_to_edge_transform_and_lower=True,