Address PR comments

Author: KyleHerndon

sharktank/sharktank/utils/iree.py

@@ -14,8 +14,6 @@
     overload,
     TYPE_CHECKING,
     Sequence,
-    Protocol,
-    runtime_checkable,
 )
 import os
 import sys
@@ -251,46 +249,13 @@ def adapt_torch_module_to_iree(
 
     iree_devices = get_iree_devices(device=device, device_count=device_count)
 
-    def load_fn(devices: list[iree.runtime.HalDevice]) -> TorchLikeIreeModule:
-        vm_module, vm_context, vm_instance = load_iree_module(
-            module_buff=vmfb_bytes,
-            devices=devices,
-            parameters_path=parameters_path,
-            tensor_parallel_size=len(devices),
-        )
-        return TorchLikeIreeModule(vm_module, vm_context, devices)
-
-    return with_iree_device_context(load_fn, iree_devices)
-
-
-@runtime_checkable
-class InferenceModule(Protocol):
-    """Protocol for inference modules (both torch and IREE).
-
-    This defines a common interface that both torch.nn.Module and
-    TorchLikeIreeModule can satisfy, allowing them to be used
-    interchangeably in inference code.
-
-    Example:
-        >>> def run_inference(model: InferenceModule, inputs):
-        ...     return model(inputs)
-        >>>
-        >>> # Works with torch modules
-        >>> torch_model = MyTorchModel()
-        >>> run_inference(torch_model, x)
-        >>>
-        >>> # Also works with IREE modules
-        >>> iree_model = adapt_torch_module_to_iree(torch_model, ...)
-        >>> run_inference(iree_model, x)
-    """
-
-    def __call__(self, *args: Any, **kwargs: Any) -> Any:
-        """Execute the module's forward pass."""
-        ...
-
-    def forward(self, *args: Any, **kwargs: Any) -> Any:
-        """Execute the module's forward pass explicitly."""
-        ...
+    vm_module, vm_context, vm_instance = load_iree_module(
+        module_buff=vmfb_bytes,
+        devices=iree_devices,
+        parameters_path=parameters_path,
+        tensor_parallel_size=len(iree_devices),
+    )
+    return TorchLikeIreeModule(vm_module, vm_context, iree_devices)
 
 
 class TorchLikeIreeModule:

sharktank/tests/utils/iree_test.py

@@ -129,12 +129,12 @@ def roundtrip(iree_devices: list[iree.runtime.HalDevice]):
 
 class SimpleModel(nn.Module):
     def forward(self, x):
-        return torch.relu(x) + 1.0
+        return x + 1
 
 
 class MultiOutputModel(nn.Module):
     def forward(self, x):
-        return torch.relu(x), torch.tanh(x)
+        return x + 1, x * 2
 
 
 class TestCompileTorchModule:
@@ -167,10 +167,10 @@ def test_compilation(self, tmp_path):
 class TestAdaptTorchModuleToIree:
     """Tests for adapt_torch_module_to_iree."""
 
-    def test_basic_loading_and_execution(self):
+    def test_basic_loading_and_execution(self, deterministic_random_seed):
         """Test that loaded module executes and produces correct output shape."""
         model = SimpleModel()
-        example_input = torch.randn(2, 32)
+        example_input = torch.randint(0, 100, (2, 32), dtype=torch.int64)
 
         iree_module = adapt_torch_module_to_iree(
             model,
@@ -182,14 +182,12 @@ def test_basic_loading_and_execution(self):
         result = iree_module.forward(example_input)
         assert isinstance(result, torch.Tensor)
         assert result.shape == (2, 32)
-        assert not torch.isnan(result).any()
 
-    def test_output_matches_torch(self):
-        """Test that IREE output matches torch output."""
-        torch.manual_seed(42)
+    def test_output_matches_torch(self, deterministic_random_seed):
+        """Test that IREE output matches torch output"""
         model = SimpleModel()
         model.eval()
-        example_input = torch.randn(2, 32)
+        example_input = torch.randint(0, 100, (2, 32), dtype=torch.int64)
 
         torch_output = model(example_input)
         iree_module = adapt_torch_module_to_iree(
@@ -200,7 +198,7 @@ def test_output_matches_torch(self):
         )
         iree_output = iree_module.forward(example_input)
 
-        torch.testing.assert_close(iree_output, torch_output, rtol=1e-4, atol=1e-4)
+        assert torch.equal(iree_output, torch_output)
 
     def test_multi_output_model(self):
         """Test model with multiple outputs."""
@@ -224,10 +222,10 @@ def test_multi_output_model(self):
 class TestOneshotCompileAndRun:
     """Tests for oneshot_iree_run."""
 
-    def test_basic_oneshot(self):
+    def test_basic_oneshot(self, deterministic_random_seed):
         """Test basic one-shot execution."""
         model = SimpleModel()
-        example_input = torch.randn(2, 32)
+        example_input = torch.randint(0, 100, (2, 32), dtype=torch.int64)
 
         result = oneshot_iree_run(
             model,
@@ -237,14 +235,12 @@ def test_basic_oneshot(self):
         )
         assert isinstance(result, torch.Tensor)
         assert result.shape == (2, 32)
-        assert not torch.isnan(result).any()
 
-    def test_oneshot_matches_torch(self):
+    def test_oneshot_matches_torch(self, deterministic_random_seed):
         """Test that one-shot execution matches torch."""
-        torch.manual_seed(42)
         model = SimpleModel()
         model.eval()
-        example_input = torch.randn(2, 32)
+        example_input = torch.randint(0, 100, (2, 32), dtype=torch.int64)
 
         torch_output = model(example_input)
         iree_output = oneshot_iree_run(
@@ -254,4 +250,5 @@ def test_oneshot_matches_torch(self):
             compile_args=COMPILE_FLAGS,
         )
 
-        torch.testing.assert_close(iree_output, torch_output, rtol=1e-4, atol=1e-4)
+        # Use exact comparison for integer arithmetic
+        assert torch.equal(iree_output, torch_output)