Fix after review

Author: irtrukhina

backends/nxp/tests/dataset_creator.py

@@ -348,7 +348,7 @@ class LinearRampDatasetCreator(DatasetCreator):
     equal to `high`, with increments depending on the total number of elements.
     """
 
-    def __init__(self, num_samples=2, low=0.0, high=1.0):
+    def __init__(self, num_samples=2, low=-1.0, high=1.0):
         self._num_samples = num_samples
         self.low = low
         self.high = high

backends/nxp/tests/ir/converter/node_converter/test_log_converter.py

@@ -11,9 +11,12 @@
 
 from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier
 from executorch.backends.nxp.tests.nsys_testing import lower_run_compare
-from executorch.backends.nxp.tests.ops_aliases import Clamp, Convolution, Log, Relu
+from executorch.backends.nxp.tests.ops_aliases import Log
 from executorch.backends.nxp.tests.use_qat import *  # noqa F403
-from executorch.backends.nxp.tests.dataset_creator import LinearRampDatasetCreator
+from executorch.backends.nxp.tests.dataset_creator import (
+    LinearRampDatasetCreator,
+    RandomDatasetCreator,
+)
 
 
 @pytest.fixture(autouse=True)
@@ -31,40 +34,12 @@ def forward(self, x):
         return torch.log(x)
 
 
-class ConvBlocksWithLogModule(torch.nn.Module):
-    def __init__(self, conv_in_channels: int = 3):
-        super().__init__()
-        self.block1 = torch.nn.Sequential(
-            torch.nn.Conv2d(
-                in_channels=conv_in_channels,
-                out_channels=3,
-                kernel_size=(2, 2),
-                stride=(2, 2),
-            ),
-            torch.nn.ReLU(),
-        )
-        self.block2 = torch.nn.Sequential(
-            torch.nn.Conv2d(
-                in_channels=conv_in_channels,
-                out_channels=10,
-                kernel_size=(2, 2),
-                stride=(2, 2),
-            ),
-            torch.nn.ReLU(),
-        )
-
-    def forward(self, x):
-        x = self.block1(x)
-        x = x.clamp_min(1e-6).log()
-        return self.block2(x)
-
-
 class TestLog:
     def test__basic_nsys_inference(self, mocker):
-        # Use 256 elements so that, after quantization to uint8, the input can
-        # cover the full discrete range [0, 255].
+        # Use 256 elements so that, after quantization to int8, the input can
+        # cover the full discrete range [-128, 127].
         # The dataset is generated as a linear float ramp and later quantized,
-        # which effectively exercises all uint8 values.
+        # which effectively exercises all int8 values.
         input_shape = (256,)
         model = LogModule()
         graph_verifier = DetailedGraphVerifier(
@@ -74,23 +49,26 @@ def test__basic_nsys_inference(self, mocker):
             model,
             input_shape,
             graph_verifier,
-            dataset_creator=LinearRampDatasetCreator(),
+            dataset_creator=LinearRampDatasetCreator(low=0.0, high=1.0),
         )
 
-    def test__basic_nsys_inference__with_conv(self, mocker):
-        input_shape = (2, 3, 6, 7)
-        in_channels = input_shape[1]
-        model = ConvBlocksWithLogModule(conv_in_channels=in_channels)
-
-        # `clamp` and one `relu` ends up in the same delegated partition as `log`
+    @pytest.mark.parametrize(
+        "input_shape",
+        [
+            pytest.param((17, 2), id="2D"),
+            pytest.param((1, 3, 10), id="3D"),
+            pytest.param((1, 3, 16, 16), id="4D"),
+        ],
+    )
+    def test__basic_nsys_inference__qat(self, mocker, input_shape, use_qat):
+        model = LogModule()
         graph_verifier = DetailedGraphVerifier(
-            mocker,
-            expected_delegated_ops={Log: 1, Relu: 1, Clamp: 1},
-            expected_non_delegated_ops={Relu: 1, Convolution: 2},
+            mocker, expected_delegated_ops={Log: 1}, expected_non_delegated_ops={}
         )
-
         lower_run_compare(
             model,
             input_shape,
             graph_verifier,
+            dataset_creator=RandomDatasetCreator(low=1.0, high=10.0),
+            use_qat=use_qat,
         )