Anandhu P Shaji

Hunting Compiler Bugs and Numerical Traps in TensorFlow

· Anandhu P Shaji · Back to Code

Machine learning engineers tend to treat framework primitives as immutable mathematical truths. When we write y = tf.math.reciprocal(tf.math.reciprocal(x)) or call tf.linalg.det(matrix) during a backward pass, we assume that the underlying graph optimizer, intermediate representation compiler, and C++ kernels will evaluate the expression without altering IEEE 754 semantics or silently dumping core.

Over the past year, I contributed 24 pull requests to tensorflow/tensorflow (22 merged, 2 approved). Rather than touching surface-level documentation or cosmetic wrappers, my focus was on root-causing low-level failures across Grappler (TensorFlow's computation graph optimizer), XLA (Accelerated Linear Algebra compiler), and Core numerical routines.

Here are the stories behind several of those bugs: what went wrong, why standard unit tests missed them, and how we tracked them down to root cause.

Story 1: The Case of the Unsound Involution (1 / (1 / x) == x)

Pull Requests: #123195 and #123282 (Merged into Grappler)

One of Grappler's core responsibilities is algebraic simplification: rewriting computation subgraphs into mathematically equivalent, cheaper forms before kernel execution. For example, replacing x + 0 with x or simplifying double negations -(-x) to x saves memory bandwidth and kernel launches.

In pure real analysis, reciprocal is an involution:

f(f(x)) = x, where f(x) = 1 / x

Grappler's algebraic optimizer had generalized involution elimination across unary ops, including Reciprocal. But IEEE 754 floating-point numbers do not behave like the real field ℝ. Consider what happens at floating-point boundaries: 

# Without Grappler simplification (Exact IEEE 754):
x = tf.constant(0.0, dtype=tf.float32)
r1 = tf.math.reciprocal(x)    # 1.0 / 0.0 -> +inf
r2 = tf.math.reciprocal(r1)   # 1.0 / +inf -> 0.0 (preserved)

# But what about subnormals or values near float32 max (~3.4e38)?
x = tf.constant(1e-40, dtype=tf.float32)  # subnormal
r1 = tf.math.reciprocal(x)               # overflows to +inf!
r2 = tf.math.reciprocal(r1)               # evaluates to 0.0 != 1e-40!

When Grappler aggressively eliminated Reciprocal(Reciprocal(x)) → x, it bypassed the intermediate overflow to +inf. Worse, if x = -0.0, strict floating-point division preserves signed zeroes (1 / -0.0 = -inf, 1 / -inf = -0.0), whereas certain hardware targets flushed signs during graph rewrites.

The fix required carving out Reciprocal from unsafe involution rules in Grappler's C++ optimizer (tensorflow/core/grappler/optimizers/arithmetic_optimizer.cc), ensuring that algebraic rewrites never compromise IEEE 754 compliance for models with extreme dynamic range.

Story 2: The Silent Compiler Segfault in XLA (TensorArray.unstack)

Pull Requests: #124929 and #124930 (Merged into XLA)

XLA compiles TensorFlow subgraphs into specialized machine code using LLVM or GPU PTX. Because XLA performs ahead-of-time (AOT) and just-in-time (JIT) shape inference, any mismatch between expected rank and actual shape can lead to undefined memory access.

We discovered that invoking TensorArray.unstack or TensorArrayScatterOp inside an XLA-compiled context with a scalar (0-D) input caused the entire Python runtime to terminate immediately with a Segmentation Fault: 

import tensorflow as tf

@tf.function(jit_compile=True)
def test_unstack():
    ta = tf.TensorArray(dtype=tf.float32, size=0, dynamic_size=True)
    # Scalar input passed into unstack:
    return ta.unstack(tf.constant(42.0))

test_unstack()  # Fatal: Segmentation fault (core dumped)

Tracking this down required running GDB attached to the Python binary. In the XLA shape translation pass: 

// Inside xla/service/shape_inference.cc:
xla::Shape element_shape = tensor_shape;
element_shape.DeleteDimension(0);  // Unchecked assumption: rank >= 1

Because tensor_shape had rank 0, DeleteDimension(0) attempted to resize an empty vector and read index 0 out of bounds. The compiler crashed before emitting any diagnostic error message.

In PR #124929 and #124930, we introduced rank validation checks at the XLA lowering boundary, translating what was an unrecoverable C++ segfault into a clean, actionable InvalidArgumentError: Input tensor must have rank >= 1 to unstack.

Story 3: Singular Matrices and Exploding Gradients in tf.linalg.det

Pull Request: #122823 (Merged into Core Math / Gradients)

In deep learning, matrix determinants appear in Gaussian process log-likelihoods, normalizing flows, and entropy regularization. Automatic differentiation of the matrix determinant relies on Jacobi's formula:

d(det(A)) / dA = det(A) · (A-1)T

Mathematically, this identity holds even when det(A) = 0, provided one takes the limit using the matrix adjugate: adj(A)T.

In practice, however, TensorFlow computed the gradient by calculating A-1 numerically using LU decomposition and then multiplying by det(A): 

# Backpropagation through a singular matrix:
A = tf.Variable([[1.0, 2.0], [2.0, 4.0]])  # det(A) = 0.0
with tf.GradientTape() as tape:
    d = tf.linalg.det(A)
grad = tape.gradient(d, A)

# Expected gradient: [[4.0, -2.0], [-2.0, 1.0]] (adjugate transpose)
# Actual result in TensorFlow: [[NaN, NaN], [NaN, NaN]]

Because A is singular, the numerical inverse contained inf and NaN entries. The subsequent multiplication 0.0 * inf produced NaN, silently destroying the optimizer's momentum state in any training loop where matrices passed through a singular boundary.

We overhauled the gradient kernel to detect singular conditions and gracefully handle non-invertible matrices during the backward pass, preventing catastrophic NaN contamination in downstream gradient updates.

Story 4: The IEEE 754 Boundary Traps (-inf, igamma, and logsumexp)

Pull Requests: #123862, #124927, #115554

Floating-point corner cases are notorious for slipping past standard random-value unit tests. Three PRs addressed subtle edge cases that only surfaced under extreme inputs:

Lessons from 24 Pull Requests

Working on a library of TensorFlow's scale - with millions of lines of C++, Python, MLIR, and CUDA - reinforced several core engineering principles:

  1. Compilers must respect hardware floating-point standards: High-level algebraic laws (commutativity, associativity, invertibility) do not hold in finite-precision IEEE 754. Optimization passes must verify that algebraic rewrites do not alter behavior on subnormals, infinities, or signed zeroes.
  2. Crash early with informative diagnostics: Unchecked assumptions in internal compiler passes (such as vector indexing in shape inference) turn minor user mistakes into fatal segmentation faults. Validating invariants at the boundary protects the entire runtime.
  3. Differential testing beats random testing: Standard test suites with uniform random numbers almost never hit singular matrices, negative zeroes, or extreme ratios. Generating adversarial edge cases is essential for rock-solid numerical libraries.

All 24 pull requests, including test cases, code reviews, and discussions, can be viewed directly on GitHub (@Cyrax321). For an overview of my systems and open-source projects, visit the Code page.