Fix (brevitas_examples/quantizers): correct stats for dynamic quants (#1445)

Author: Giuseppe5

src/brevitas_examples/common/generative/quantizers.py

@@ -11,6 +11,7 @@
 from brevitas.core.scaling.runtime import RuntimeDynamicGroupStatsScaling
 from brevitas.core.stats import AbsMinMax
 from brevitas.core.stats import NegativeMinOrZero
+from brevitas.core.stats import StatsOp
 from brevitas.core.stats.stats_op import HalfQuadraticOptimizerZeroPoint
 from brevitas.core.stats.stats_wrapper import SCALAR_SHAPE
 from brevitas.core.zero_point import RuntimeDynamicGroupZeroPoint
@@ -72,7 +73,7 @@ class Int8DynamicActPerTensorFloat(DynamicActProxyMixin, Int8ActPerTensorFloat):
     """
     scaling_impl = RuntimeDynamicStatsScaling
     scaling_stats_input_view_shape_impl = OverTensorView
-    scaling_stats_op = 'min_max'
+    scaling_stats_op = StatsOp.MAX
     dynamic_scaling_broadcastable_fn = lambda x, shape: x.view(SCALAR_SHAPE)
 
 
@@ -82,7 +83,7 @@ class Fp8e4m3FNUZDynamicActPerTensorFloat(Fp8e4m3FNUZActPerTensorFloat):
     """
     scaling_impl = RuntimeDynamicStatsScaling
     scaling_stats_input_view_shape_impl = OverTensorView
-    scaling_stats_op = 'min_max'
+    scaling_stats_op = StatsOp.MAX
     dynamic_scaling_broadcastable_fn = lambda x, shape: x.view(SCALAR_SHAPE)
 
 
@@ -92,7 +93,7 @@ class Int8DynamicActPerRowFloat(DynamicActProxyMixin, Int8ActPerTensorFloat):
     """
     scaling_impl = RuntimeDynamicStatsScaling
     scaling_stats_input_view_shape_impl = OverOutputFeaturesView
-    scaling_stats_op = 'min_max'
+    scaling_stats_op = StatsOp.MAX
     scaling_per_output_channel = True
 
 
@@ -107,7 +108,7 @@ class Int8DynamicActPerGroupFloat(DynamicActProxyMixin, Int8ActPerTensorFloat):
     """
     proxy_class = GroupwiseActQuantProxyFromInjector
     scaling_impl = RuntimeDynamicGroupStatsScaling
-    scaling_stats_op = 'min_max'
+    scaling_stats_op = StatsOp.MAX
     scaling_per_output_type = ScalingPerOutputType.GROUP
 
 
@@ -117,7 +118,7 @@ class ShiftedUint8DynamicActPerGroupFloat(DynamicActProxyMixin, ShiftedUint8ActP
     """
     proxy_class = GroupwiseActQuantProxyFromInjector
     scaling_impl = RuntimeDynamicGroupStatsScaling
-    scaling_stats_op = 'min_max'
+    scaling_stats_op = StatsOp.MIN_MAX
     scaling_per_output_type = ScalingPerOutputType.GROUP
     zero_point_impl = RuntimeDynamicGroupZeroPoint
     zero_point_stats_impl = NegativeMinOrZero
@@ -129,7 +130,7 @@ class ShiftedUint8DynamicActPerTensorFloat(DynamicActProxyMixin, ShiftedUint8Act
     """
     scaling_impl = RuntimeDynamicStatsScaling
     scaling_stats_input_view_shape_impl = OverTensorView
-    scaling_stats_op = 'min_max'
+    scaling_stats_op = StatsOp.MIN_MAX
     zero_point_impl = RuntimeDynamicStatsZeroPoint
     zero_point_stats_impl = NegativeMinOrZero
     dynamic_scaling_broadcastable_fn = lambda x, shape: x.view(SCALAR_SHAPE)
@@ -141,7 +142,7 @@ class ShiftedUint8DynamicActPerRowFloat(DynamicActProxyMixin, ShiftedUint8ActPer
     """
     scaling_impl = RuntimeDynamicStatsScaling
     scaling_stats_input_view_shape_impl = OverOutputFeaturesView
-    scaling_stats_op = 'min_max'
+    scaling_stats_op = StatsOp.MIN_MAX
     scaling_per_output_channel = True
     zero_point_impl = RuntimeDynamicStatsZeroPoint
     zero_point_stats_impl = NegativeMinOrZero
@@ -154,7 +155,7 @@ class Fp8e4m3DynamicActPerGroupFloat(DynamicActProxyMixin, Fp8e4m3ActPerTensorFl
     proxy_class = GroupwiseActFloatQuantProxyFromInjector
     scaling_impl = RuntimeDynamicGroupStatsScaling
     scaling_per_output_type = ScalingPerOutputType.GROUP
-    scaling_stats_op = 'min_max'
+    scaling_stats_op = StatsOp.MAX
 
 
 class FP8e4m3OCPDynamicActPerRowFixedPoint(Fp8e4m3OCPActPerTensorFloat):
@@ -163,7 +164,7 @@ class FP8e4m3OCPDynamicActPerRowFixedPoint(Fp8e4m3OCPActPerTensorFloat):
     """
     scaling_impl = RuntimeDynamicStatsScaling
     scaling_stats_input_view_shape_impl = OverOutputFeaturesView
-    scaling_stats_op = 'min_max'
+    scaling_stats_op = StatsOp.MAX
     scaling_per_output_channel = True
     restrict_scaling_type = RestrictValueType.POWER_OF_TWO
     restrict_value_float_to_int_impl = FloorSte
@@ -173,7 +174,7 @@ class FP8e4m3OCPDynamicActPerRowFixedPoint(Fp8e4m3OCPActPerTensorFloat):
 class FP8e4m3OCPDynamicActPerRowFloat(Fp8e4m3OCPActPerTensorFloat):
     scaling_impl = RuntimeDynamicStatsScaling
     scaling_stats_input_view_shape_impl = OverOutputFeaturesView
-    scaling_stats_op = 'min_max'
+    scaling_stats_op = StatsOp.MAX
     scaling_per_output_channel = True
     proxy_class = DynamicActFloatQuantProxyFromInjector
 
@@ -185,7 +186,7 @@ class Fp8e4m3OCPDynamicActPerGroupFloat(DynamicActProxyMixin, Fp8e4m3OCPActPerTe
     proxy_class = GroupwiseActFloatQuantProxyFromInjector
     scaling_impl = RuntimeDynamicGroupStatsScaling
     scaling_per_output_type = ScalingPerOutputType.GROUP
-    scaling_stats_op = 'min_max'
+    scaling_stats_op = StatsOp.MAX
 
 
 class Fp8e4m3OCPWeightSymmetricGroupQuant(Fp8e4m3OCPWeightPerChannelFloat):
@@ -210,7 +211,7 @@ class Fp8e4m3OCPWeightPerChannelFloatMSE(MSESymmetricScale, Fp8e4m3OCPWeightPerC
 class FP8e4m3FNUZDynamicActPerRowFloat(Fp8e4m3FNUZActPerTensorFloat):
     scaling_impl = RuntimeDynamicStatsScaling
     scaling_stats_input_view_shape_impl = OverOutputFeaturesView
-    scaling_stats_op = 'min_max'
+    scaling_stats_op = StatsOp.MAX
     scaling_per_output_channel = True
     proxy_class = DynamicActFloatQuantProxyFromInjector