INT4 ONNX Version Fix: Code Quality Improvements

Signed-off-by: Hrishith Thadicherla <[email protected]>

Author: hthadicherla

modelopt/onnx/quantization/gs_patching.py

@@ -69,8 +69,8 @@ def _export_tensor_proto(tensor: gs.Constant) -> onnx.TensorProto:
 
         vals = tensor.values
         if _onnx_supports_int4() and dtype in [onnx.TensorProto.INT4, onnx.TensorProto.UINT4]:
-            signed = dtype == onnx.TensorProto.INT4      
-            if(signed):
+            signed = dtype == onnx.TensorProto.INT4
+            if signed:
                 vals = pack_float32_to_4bit_cpp_based(tensor.values, signed=signed).astype(np.int8)
             else:
                 vals = pack_float32_to_4bit_cpp_based(tensor.values, signed=signed).astype(np.uint8)

modelopt/onnx/quantization/int4.py

@@ -98,11 +98,13 @@
 # supported and working
 CLIP_MIN = 1e-5
 
+
 def safe_cupy_array(tensor):
     """Convert ml_dtypes.int4 tensor to numpy.int8 for CuPy compatibility.
 
     In ONNX 1.19, int4 tensors use ml_dtypes.int4 which CuPy doesn't support.
     This function converts them to regular numpy.int8 while preserving values.
+
     Args:
         tensor: numpy array that may have ml_dtypes.int4 dtype
     Returns:
@@ -111,12 +113,12 @@ def safe_cupy_array(tensor):
     """
     try:
         import ml_dtypes
-        
-        if hasattr(tensor, 'dtype') and tensor.dtype == ml_dtypes.int4:
+
+        if hasattr(tensor, "dtype") and tensor.dtype == ml_dtypes.int4:
             return np.asarray(tensor.astype(numpy.int8))
     except ImportError:
         pass
-    
+
     return np.asarray(tensor)