coderabbitai fixes

Signed-off-by: Riyad Islam <[email protected]>

Author: i-riyad

modelopt/onnx/quantization/qdq_utils.py

@@ -527,11 +527,13 @@ def _get_successive_consumers(
         raise ValueError(f"Invalid consumer for {node.name}")
 
     quantized_node = tensor_consumers.get(dq_node.output[0], [None])[0]
-    if quantized_node.op_type == "Cast":
-        quantized_node = tensor_consumers.get(quantized_node.output[0], [None])[0]
-
     if not quantized_node:
         raise ValueError(f"No consumer found for {dq_node.name}")
+    if quantized_node.op_type == "Cast":
+        next_node = tensor_consumers.get(quantized_node.output[0], [None])[0]
+        if not next_node:
+            raise ValueError(f"No consumer found after Cast for {quantized_node.name}")
+        quantized_node = next_node
 
     return dq_node, quantized_node
 

modelopt/torch/_deploy/_runtime/tensorrt/tensorrt_utils.py

@@ -131,11 +131,6 @@ def get_output_shapes(
     return output_shapes
 
 
-def validate_precision(precision: str) -> bool:
-    """Returns whether an input precision is in supported set."""
-    return precision in [TRTMode.FLOAT32, TRTMode.FLOAT16, TRTMode.INT8]
-
-
 def calib_data_generator(onnx_bytes: bytes, input_tensors: list[np.ndarray]):
     """The calibation data generator that yields calibration feed_dict to tensorrt."""
     input_names = get_onnx_input_names(onnx.load_from_string(onnx_bytes))

tests/_test_utils/onnx_quantization/utils.py

@@ -20,11 +20,11 @@ def _assert_nodes_are_quantized(nodes):
     for node in nodes:
         for inp_idx, inp in enumerate(node.inputs):
             if isinstance(inp, gs.Variable):
-                qnode = node
-                # After quantization, the quantized node can be casted
-                if qnode.i(inp_idx).op == "Cast":
-                    qnode = qnode.i(inp_idx)
-                assert qnode.i(inp_idx).op == "DequantizeLinear", (
-                    f"Input '{inp.name}' of node '{qnode.name}' is not quantized but should be!"
+                producer = node.i(inp_idx)
+                # Quantized path may include a Cast right after DQ
+                if producer and producer.op == "Cast":
+                    producer = producer.i(0)
+                assert producer and producer.op == "DequantizeLinear", (
+                    f"Input '{inp.name}' of node '{node.name}' is not quantized but should be!"
                 )
     return True