Fix for models (#2555)

* Handle zero dimensions corner cases for the torch tenson assign

* expose calibration_op_group_size arg

Author: jakesabathia2

coremltools/converters/mil/frontend/torch/ops.py

@@ -4836,17 +4836,23 @@ def _internal_op_tensor_inplace_fill(context, node):
 
     data, update_values = promote_input_dtypes([data, update_values])
 
-    updated_x = _translate_torch_tensor_assign(
-        x=data,
-        updates=update_values,
-        begin=begin,
-        end=end,
-        stride=stride,
-        begin_mask=begin_mask,
-        end_mask=end_mask,
-        squeeze_mask=squeeze_mask,
-        name=node.name,
-    )
+    if 0 in update_values.shape:
+        # if the update_values contains zero dimensions, this will be a noop
+        updated_x = mb.identity(x=data, name=node.name)
+    else:
+        # otherwise we translate it into a tensor assignment op
+        updated_x = _translate_torch_tensor_assign(
+            x=data,
+            updates=update_values,
+            begin=begin,
+            end=end,
+            stride=stride,
+            begin_mask=begin_mask,
+            end_mask=end_mask,
+            squeeze_mask=squeeze_mask,
+            name=node.name,
+        )
+
     context.add(updated_x)
 
 

coremltools/converters/mil/frontend/torch/test/test_torch_ops.py

@@ -9523,6 +9523,31 @@ def forward(self, x):
             assert "slice_update" in get_op_types_in_program(prog)
 
 
+    @pytest.mark.parametrize(
+        "compute_unit, backend, minimum_deployment_target",
+        itertools.product(
+            compute_units,
+            backends,
+            [None, ct.target.iOS18],
+        ),
+    )
+    def test_tensor_assign_no_op(self, compute_unit, backend, minimum_deployment_target):
+        # if the resulting update has `0` dimensions, and it is a noop
+        class TensorAssignModel(torch.nn.Module):
+            def forward(self, x):
+                x[:, -8:0, :] = 0
+                return x + 1
+
+        shape = (2, 10, 3)
+        model = TensorAssignModel()
+        res = self.run_compare_torch(
+            shape,
+            model,
+            backend=backend,
+            compute_unit=compute_unit,
+        )
+
+
 class TestSelectScatter(TorchBaseTest):
     @pytest.mark.parametrize(
         "compute_unit, backend, frontend, minimum_deployment_target, input_shape, dynamic",

coremltools/optimize/coreml/_post_training_quantization.py

@@ -601,7 +601,11 @@ def linear_quantize_activations(
         "compression::insert_prefix_quantize_dequantize_pair"
     ]
     insert_prefix_quantize_dequantize_pair.set_options([PassOption("config", config)])
-    activation_stats = _get_activation_calibration_stats(mlmodel, sample_data)
+    activation_stats = _get_activation_calibration_stats(
+        mlmodel, 
+        sample_data,
+        calibration_op_group_size,
+    )
     insert_prefix_quantize_dequantize_pair.set_options(
         [PassOption("activation_stats", activation_stats)]
     )

coremltools/optimize/coreml/experimental/test_post_training_quantization.py

@@ -5,11 +5,15 @@
 
 import numpy as np
 import torch
+from io import StringIO
+import sys
 
 import coremltools as ct
 from coremltools.optimize.coreml.experimental._post_training_quantization import (
     _get_activation_calibration_stats,
 )
+from coremltools.test.optimize.coreml.test_passes import TestCompressionPasses
+import coremltools.optimize as cto
 
 
 class TestActivationQuantization:
@@ -72,6 +76,45 @@ def forward(self, img):  # convert + flatten
 
 
 class TestGetActivationStats(TestActivationQuantization):
+
+    def test_activation_quantization(self):
+        """
+        Test the usage of linear_quantize_activations.
+        """
+        sample_data = []
+        for _ in range(3):
+            input_data = np.random.rand(5, 10, 4, 4)
+            sample_data.append({"data": input_data})
+
+        mlmodel = self._get_test_mlmodel_conv_relu()
+        activation_quant_config = cto.coreml.OptimizationConfig(
+            global_config=cto.coreml.experimental.OpActivationLinearQuantizerConfig(
+                mode="linear_symmetric", weight_threshold=10
+            )
+        )
+
+        def _run_quantization_with_group_size(group_size, expected_batch_size):
+            buffer = StringIO()
+            original_stderr = sys.stderr
+            sys.stderr = buffer
+            mlmodel_activation_quantized = cto.coreml.experimental.linear_quantize_activations(
+                mlmodel,
+                activation_quant_config,
+                sample_data,
+                calibration_op_group_size=group_size,
+            )
+            sys.stderr = original_stderr
+            output = buffer.getvalue()
+            assert f"tensors batch-by-batch: {expected_batch_size} batches" in output
+        
+        # when setting group size to -1, all intermediate outputs are in the same batch,
+        # hence we will only get 1 batch
+        _run_quantization_with_group_size(-1, 1)
+        # when setting group size to 1, all intermediate outputs are split into different batches,
+        # hence there will be 3 batches
+        _run_quantization_with_group_size(1, 3)
+
+
     def test_get_activation_calibration_stats_basic(self):
         """
         Calibration a floating point model with sample data.
@@ -113,7 +156,7 @@ def test_get_activation_calibration_stats_concat_surrounding_ops(self):
             sample_data.append({"data_0": input_data})
 
         # Loading a floating point mlmodel
-        mlmodel = self._get_test_mlmodel_conv_concat()
+        mlmodel = TestCompressionPasses._get_test_mlmodel_conv_concat()
 
         activation_stats = _get_activation_calibration_stats(mlmodel, sample_data)