[Mlir-commits] [mlir] [mlir][quant] Skip non-finite scales in fake quant bounding (PR #209985)

[Jesse Rosenstock]

https://github.com/jmr created https://github.com/llvm/llvm-project/pull/209985

Skip quantization if the dynamically generated scale is not finite. This occurs frequently in quantized JAX models where padding values bounded at -jnp.inf trigger overflow during fake quantization bounding.

Note: FakeQuantSupport.cpp acts as a legacy translation utility whose primary consumers live downstream, and there are currently no in-tree tests covering it. Therefore, no tests were added for this fix.

Assisted-by: Gemini

>From 93cb7b224d47dc732ca3757adfe52818a4fe535f Mon Sep 17 00:00:00 2001
From: Jesse Rosenstock <jmr at google.com>
Date: Thu, 16 Jul 2026 08:28:54 +0200
Subject: [PATCH] [mlir][quant] Skip non-finite scales in fake quant bounding

Skip quantization if the dynamically generated scale is not finite.
This occurs frequently in quantized JAX models where padding values
bounded at -jnp.inf trigger overflow during fake quantization bounding.

Note: FakeQuantSupport.cpp acts as a legacy translation utility whose
primary consumers live downstream, and there are currently no in-tree
tests covering it. Therefore, no tests were added for this fix.

Assisted-by: Gemini
---
 mlir/lib/Dialect/Quant/Utils/FakeQuantSupport.cpp | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/mlir/lib/Dialect/Quant/Utils/FakeQuantSupport.cpp b/mlir/lib/Dialect/Quant/Utils/FakeQuantSupport.cpp
index fb27640bfd278..97637ab71ff38 100644
--- a/mlir/lib/Dialect/Quant/Utils/FakeQuantSupport.cpp
+++ b/mlir/lib/Dialect/Quant/Utils/FakeQuantSupport.cpp
@@ -131,6 +131,12 @@ mlir::quant::fakeQuantAttrsToType(Location loc, unsigned numBits, double rmin,
   int64_t nudgedZeroPoint;
   getNudgedScaleAndZeroPoint(qmin, qmax, rmin, rmax, scale, nudgedZeroPoint);
 
+  // Skip quantization if the dynamically generated scale is not finite.
+  // This occurs frequently in quantized JAX models where padding values
+  // bounded at -jnp.inf trigger overflow during fake quantization bounding.
+  if (!std::isfinite(scale))
+    return nullptr;
+
   return UniformQuantizedType::getChecked(loc, flags, storageType,
                                           expressedType, scale, nudgedZeroPoint,
                                           qmin, qmax);
@@ -173,6 +179,8 @@ UniformQuantizedPerAxisType mlir::quant::fakeQuantAttrsToType(
     double scale;
     int64_t nudgedZeroPoint;
     getNudgedScaleAndZeroPoint(qmin, qmax, rmin, rmax, scale, nudgedZeroPoint);
+    if (!std::isfinite(scale))
+      return nullptr;
     scales.push_back(scale);
     zeroPoints.push_back(nudgedZeroPoint);
   }