Add round to nearest logic for numpy case

Author: nikita-savelyevv

nncf/quantization/algorithms/weight_compression/weight_lowering.py

@@ -554,6 +554,19 @@ def _calculate_float_quantized_weight(norm_weight: Tensor, mode: CompressWeights
     quantile_centers = fns.from_numpy(quantile_centers_np, backend=norm_weight.backend)
     indexes = fns.searchsorted(quantile_centers, norm_weight)
     quantiles = fns.from_numpy(quantiles_np, backend=indexes.backend)
+
+    if mode == CompressWeightsMode.E2M1:
+        # Round to the nearest even quantile
+        shifted_indexes = fns.clip(indexes + 1, 0, quantiles.size - 1)
+        left = quantiles[indexes]
+        right = quantiles[shifted_indexes]
+        dist_left = fns.abs(norm_weight - left)
+        dist_right = fns.abs(norm_weight - right)
+        choose_right = fns.logical_or(
+            dist_right < dist_left, fns.logical_and(dist_left == dist_right, (shifted_indexes + 1) % 2 == 0)
+        )
+        indexes = fns.where(choose_right, shifted_indexes, indexes)
+
     quantized_weight = quantiles[indexes]
     return quantized_weight
 

nncf/tensor/functions/__init__.py

@@ -36,6 +36,7 @@
 from nncf.tensor.functions.numeric import isempty as isempty
 from nncf.tensor.functions.numeric import item as item
 from nncf.tensor.functions.numeric import log2 as log2
+from nncf.tensor.functions.numeric import logical_and as logical_and
 from nncf.tensor.functions.numeric import logical_or as logical_or
 from nncf.tensor.functions.numeric import masked_mean as masked_mean
 from nncf.tensor.functions.numeric import masked_median as masked_median

nncf/tensor/functions/numeric.py

@@ -612,6 +612,18 @@ def logical_or(x1: Tensor, x2: Tensor) -> Tensor:
     """
 
 
+@tensor_dispatcher
+def logical_and(x1: Tensor, x2: Tensor) -> Tensor:
+    """
+    Computes the element-wise logical AND of the given input tensors.
+    Zeros are treated as False and nonzeros are treated as True.
+
+    :param x1: The input tensor.
+    :param x2: The tensor to compute and with.
+    :return: Result of elementwise and operation between input_ and other tensor.
+    """
+
+
 @tensor_dispatcher
 def masked_mean(x: Tensor, mask: Tensor, axis: T_AXIS, keepdims: bool = False) -> Tensor:
     """

nncf/tensor/functions/numpy_numeric.py

@@ -348,6 +348,11 @@ def _(x1: T_NUMPY_ARRAY, x2: T_NUMPY_ARRAY) -> T_NUMPY_ARRAY:
     return np.logical_or(x1, x2)
 
 
+@numeric.logical_and.register
+def _(x1: T_NUMPY_ARRAY, x2: T_NUMPY_ARRAY) -> T_NUMPY_ARRAY:
+    return np.logical_and(x1, x2)
+
+
 @numeric.masked_mean.register
 def _(
     x: T_NUMPY_ARRAY,

nncf/tensor/functions/tf_numeric.py

@@ -417,6 +417,12 @@ def _(x1: tf.Tensor, x2: tf.Tensor) -> tf.Tensor:
         return tf.logical_or(x1, x2)
 
 
+@numeric.logical_and.register
+def _(x1: tf.Tensor, x2: tf.Tensor) -> tf.Tensor:
+    with tf.device(x1.device):
+        return tf.logical_and(x1, x2)
+
+
 @numeric.masked_mean.register
 def _(
     x: tf.Tensor, mask: Optional[tf.Tensor], axis: Optional[Union[int, tuple[int, ...]]], keepdims: bool = False

nncf/tensor/functions/torch_numeric.py

@@ -366,6 +366,11 @@ def _(x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:
     return torch.logical_or(x1, x2)
 
 
+@numeric.logical_and.register
+def _(x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:
+    return torch.logical_and(x1, x2)
+
+
 @numeric.masked_mean.register
 def _(x: torch.Tensor, mask: Optional[torch.Tensor], axis: T_AXIS, keepdims: bool = False) -> torch.Tensor:
     if mask is None:

nncf/tensor/tensor.py

@@ -144,6 +144,9 @@ def __ifloordiv__(self, other: Union[Tensor, T_NUMBER]) -> Tensor:
         self._data //= unwrap_tensor_data(other)
         return self
 
+    def __mod__(self, other: Union[Tensor, T_NUMBER]) -> Tensor:
+        return cast(Tensor, _call_function("_binary_op_nowarn", self, other, operator.mod))
+
     def __matmul__(self, other: Union[Tensor, T_NUMBER]) -> Tensor:
         return Tensor(self.data @ unwrap_tensor_data(other))
 

tests/openvino/optimized_functions/test_compression_functions.py

@@ -151,19 +151,18 @@ def test_optimized_compression_is_disabled(weight_shape, is_disabled, quantizati
     reason="Due to a bug in CPU plugin compression models can fail at compilation on ARM CPUs. Ticket: 164135.",
 )
 @pytest.mark.parametrize("weight_shape", [WEIGHT_SHAPE], ids=[""])
-# @pytest.mark.parametrize("config", COMPRESSION_CONFIGS, ids=[str(c) for c in COMPRESSION_CONFIGS])
-@pytest.mark.parametrize("config", FP4_COMPRESSION_CONFIGS[-2:])
+@pytest.mark.parametrize("config", COMPRESSION_CONFIGS, ids=[str(c) for c in COMPRESSION_CONFIGS])
 @pytest.mark.parametrize(
     ("quantization_task", "tensor_backend"),
     [
         (QuantizationTask.Q, TensorBackend.numpy),
-        # (QuantizationTask.Q, "auto"),
+        (QuantizationTask.Q, "auto"),
         # NumPy backend should support OV tensors as inputs only for quantization task
-        # (QuantizationTask.Q, TensorBackend.ov),
-        # (QuantizationTask.Q_DQ, TensorBackend.numpy),
-        # (QuantizationTask.Q_DQ, "auto"),
-        # (QuantizationTask.Q_DQ_RQ, TensorBackend.numpy),
-        # (QuantizationTask.Q_DQ_RQ, "auto"),
+        (QuantizationTask.Q, TensorBackend.ov),
+        (QuantizationTask.Q_DQ, TensorBackend.numpy),
+        (QuantizationTask.Q_DQ, "auto"),
+        (QuantizationTask.Q_DQ_RQ, TensorBackend.numpy),
+        (QuantizationTask.Q_DQ_RQ, "auto"),
     ],
 )
 @pytest.mark.parametrize("dtype", [TensorDataType.float32, TensorDataType.float16, TensorDataType.bfloat16])
@@ -525,8 +524,8 @@ def format_list_of_floats(lst):
                 f"NumPy result: {format_list_of_floats(numpy_result.data[not_equal_mask])}\n"
             )
             if "input" in results[ComputationBackend.OV] and "input" in results[ComputationBackend.NumPy]:
-                numpy_input = results[ComputationBackend.NumPy]['input'].data
-                ov_input = results[ComputationBackend.OV]['input'].data
+                numpy_input = results[ComputationBackend.NumPy]["input"].data
+                ov_input = results[ComputationBackend.OV]["input"].data
                 np.testing.assert_allclose(numpy_input, ov_input, atol=0, rtol=0)
                 msg += f"Input values   : {format_list_of_floats(numpy_input[not_equal_mask])}\n"
                 misaligned_groups_mask = np.any(not_equal_mask, axis=-1)