Merge branch 'main' into dev/gokul/flmrife-cat-shape

Author: gokulkrishna98

tests/compression/test_compression.py

@@ -13,7 +13,6 @@
 import pytest
 import torch
 from coreai.authoring import AIProgram
-from coreai.runtime import AIModel, NDArray
 from torch import nn
 from torch.export.exported_program import ExportedProgram
 
@@ -24,7 +23,7 @@
     ActivationQuantizeModule,
 )
 
-from ..utils import TemporaryModelAsset, filecheck_pattern
+from ..utils import filecheck_pattern, validate_numerical_output
 
 # We add "./tests/coreai" path, in order to use some existing utils
 sys.path.append(str(Path(__file__).parents[2]))
@@ -58,32 +57,6 @@ async def lower_to_coreai(
     return converter.to_coreai()
 
 
-async def _validate_execution(
-    coreai_program: AIProgram,
-    torch_out: torch.Tensor,
-    atol: float = 1e-4,
-    rtol: float = 1e-4,
-    **kwargs: Any,
-) -> None:
-    """Run the Core AI program using ref kernels and match with torch output."""
-    with TemporaryModelAsset() as tempdir:
-        coreai_program.save_asset(Path(tempdir))
-        ai_model = await AIModel.load(Path(tempdir))
-        rt_func = ai_model.load_function("main")
-
-        # Wrap all kwargs with NDArray
-        nd_kwargs = {k: NDArray(data=v) for k, v in kwargs.items()}
-        coreai_outs = await rt_func(nd_kwargs)
-
-        coreai_out_np = {k: v.numpy() for k, v in coreai_outs.items()}
-        np.testing.assert_allclose(
-            torch_out.detach().numpy(),
-            next(iter(coreai_out_np.values())),  # type: ignore[arg-type]
-            rtol=rtol,
-            atol=atol,
-        )
-
-
 @pytest.mark.parametrize(
     "nbits",
     [4, 8],
@@ -164,9 +137,11 @@ def forward(self, input_tensor: torch.Tensor) -> torch.Tensor:
         exported_program,
     )
 
-    await _validate_execution(
-        coreai_program,
-        torch_output,
+    await validate_numerical_output(
+        coreai_program=coreai_program,
+        torch_out=torch_output,
+        atol=1e-4,
+        rtol=1e-4,
         input_tensor=input_tensor,
     )
 
@@ -233,9 +208,11 @@ def forward(self, input_tensor: torch.Tensor) -> torch.Tensor:
         exported_program,
     )
 
-    await _validate_execution(
-        coreai_program,
-        torch_output,
+    await validate_numerical_output(
+        coreai_program=coreai_program,
+        torch_out=torch_output,
+        atol=1e-4,
+        rtol=1e-4,
         input_tensor=input_tensor,
     )
 
@@ -284,9 +261,11 @@ def forward(self, input_tensor: torch.Tensor) -> torch.Tensor:
         exported_program,
     )
 
-    await _validate_execution(
-        coreai_program,
-        torch_output,
+    await validate_numerical_output(
+        coreai_program=coreai_program,
+        torch_out=torch_output,
+        atol=1e-4,
+        rtol=1e-4,
         input_tensor=input_tensor,
     )
 
@@ -372,9 +351,11 @@ def forward(self, input_tensor: torch.Tensor) -> torch.Tensor:
         exported_program,
     )
 
-    await _validate_execution(
-        coreai_program,
-        torch_output,
+    await validate_numerical_output(
+        coreai_program=coreai_program,
+        torch_out=torch_output,
+        atol=1e-4,
+        rtol=1e-4,
         input_tensor=input_tensor,
     )
 
@@ -442,9 +423,11 @@ def forward(self, input_tensor: torch.Tensor) -> torch.Tensor:
         exported_program,
     )
 
-    await _validate_execution(
-        coreai_program,
-        torch_output,
+    await validate_numerical_output(
+        coreai_program=coreai_program,
+        torch_out=torch_output,
+        atol=1e-4,
+        rtol=1e-4,
         input_tensor=input_tensor,
     )
 
@@ -541,9 +524,11 @@ def forward(self, input_tensor: torch.Tensor) -> torch.Tensor:
         exported_program,
     )
 
-    await _validate_execution(
-        coreai_program,
-        torch_output,
+    await validate_numerical_output(
+        coreai_program=coreai_program,
+        torch_out=torch_output,
+        atol=1e-4,
+        rtol=1e-4,
         input_tensor=input_tensor,
     )
 
@@ -619,9 +604,11 @@ def forward(self, input_tensor: torch.Tensor) -> torch.Tensor:
         exported_program,
     )
 
-    await _validate_execution(
-        coreai_program,
-        torch_output,
+    await validate_numerical_output(
+        coreai_program=coreai_program,
+        torch_out=torch_output,
+        atol=1e-4,
+        rtol=1e-4,
         input_tensor=input_tensor,
     )
 
@@ -808,9 +795,11 @@ def forward(self, input_tensor: torch.Tensor) -> torch.Tensor:
 
         coreai_program = await lower_to_coreai(exported_program)
 
-        await _validate_execution(
-            coreai_program,
-            torch_output,
+        await validate_numerical_output(
+            coreai_program=coreai_program,
+            torch_out=torch_output,
+            atol=1e-4,
+            rtol=1e-4,
             input_tensor=input_tensor,
         )
 
@@ -954,9 +943,11 @@ def forward(self, input_tensor: torch.Tensor) -> torch.Tensor:
             exported_program,
         )
 
-        await _validate_execution(
-            coreai_program,
-            torch_output,
+        await validate_numerical_output(
+            coreai_program=coreai_program,
+            torch_out=torch_output,
+            atol=1e-4,
+            rtol=1e-4,
             input_tensor=input_tensor,
         )
 
@@ -1119,4 +1110,10 @@ async def test_numerical(
             exported_program,
         )
 
-        await _validate_execution(coreai_program, torch_out, input_tensor=input_tensor)
+        await validate_numerical_output(
+            coreai_program=coreai_program,
+            torch_out=torch_out,
+            atol=1e-4,
+            rtol=1e-4,
+            input_tensor=input_tensor,
+        )

tests/debugging/test_debug_info.py

@@ -5,6 +5,7 @@
 
 """Tests for debug_infos from CompiledLibrary and AIModel."""
 
+import sys
 import tempfile
 from pathlib import Path
 
@@ -67,6 +68,10 @@ def _verify_debug_info_record(record: DebugInfoRecord) -> None:
         ]
 
 
+@pytest.mark.skipif(
+    sys.platform != "darwin",
+    reason="Requires loading a runtime asset (AIModel.load); only supported on macOS",
+)
 @pytest.mark.asyncio
 async def test_compiled_library_debug_infos(
     simple_coreai_program: AIProgram,
@@ -91,6 +96,10 @@ async def test_compiled_library_debug_infos(
         _verify_debug_info_record(debug_info_records[0])
 
 
+@pytest.mark.skipif(
+    sys.platform != "darwin",
+    reason="Requires loading a runtime asset (AIModel.load); only supported on macOS",
+)
 @pytest.mark.asyncio
 async def test_aimodel_debug_infos(
     simple_coreai_program: AIProgram,

tests/debugging/test_inspector.py

@@ -5,6 +5,7 @@
 
 """Test inspector implementations."""
 
+import sys
 import tempfile
 from pathlib import Path
 
@@ -105,6 +106,10 @@ async def test_caching_inspector() -> None:
     )
 
 
+@pytest.mark.skipif(
+    sys.platform != "darwin",
+    reason="Requires loading a runtime asset (AIModel.load); only supported on macOS",
+)
 @pytest.mark.asyncio
 async def test_coreai_inspector(simple_coreai_program: AIProgram) -> None:
     """Test _CoreAIInspector with a deployed model."""

tests/debugging/test_validator.py

@@ -5,6 +5,7 @@
 
 """Tests for validator with hierarchical graphs."""
 
+import sys
 from typing import Any
 
 import numpy as np
@@ -406,6 +407,10 @@ async def _create_coreai_program_from_model(
     return coreai_program
 
 
+@pytest.mark.skipif(
+    sys.platform != "darwin",
+    reason="Requires loading a runtime asset (AIModel.load); only supported on macOS",
+)
 @pytest.mark.parametrize(
     "nan_branch",
     [