A finite loss does not guarantee finite gradients. Check where the mask is applied. Dividing an entire batch and selecting valid results afterward changes which values reach the loss; it does not automatically remove an earlier division by zero from the computational path. Repair the domain before the dangerous operation instead of only making the final output look clean.

Why can gradients still become NaN after torch.where or a loss mask produces normal-looking values? Two array elements separate forward selection from derivative arithmetic. They also let us compare filtering inputs before division with protecting operands while retaining a dense shape. 中文版.

Valid selected values do not validate every earlier operation

The PyTorch 2.8 autograd notes give a direct counterexample: dividing by zero and excluding the invalid output through indexing can still produce NaN during backward. They recommend masking before division. This is documented framework behavior, not a model-backward experiment executed for this article.

Reduce the example to x=(1,3), d=(0,2), and a loss using only the second, nonzero-denominator entry. Full-array division produces (inf,1.5); selecting the valid entry produces loss 1.5. The loss and selected output look normal. The first division provides different evidence about the path already evaluated.

With an invalid output replaced by zero, the upstream coefficients entering division can be written as (0,1). A local derivative path with respect to x involves dividing those coefficients by d, leaving 0/0 in the first entry. The argument that multiplying by zero removes a contribution cannot turn undefined arithmetic into an ordinary real zero. Avoiding an undefined path is a different requirement from hiding its forward value.

where selects values; it is not a lazy Python branch

The torch.where definition selects elements from its input and alternative according to a condition. In an ordinary Python call where(mask,x/d,0), the x/d argument is evaluated before the selection function receives it. The function name does not promise that division runs only at True positions.

The Python expression rules should be distinguished from conditional expressions. A scalar a if condition else b selects a branch, but a batch containing different Boolean conditions does not acquire elementwise short-circuit evaluation merely by passing arguments to a function named where. This discussion concerns ordinary eager execution, not assumptions about compiler optimization.

NumPy also provides elementwise selection, allowing an actual arithmetic experiment below. It exposes the raw infinity, the finite selected array and NaN in an explicitly written derivative expression. NumPy does not execute PyTorch automatic differentiation here. These observations must not be described as a framework backward test.

Multiplication by a mask may fail even in forward

Another approach multiplies the entire quotient by a zero/one mask. Zero times a finite value is zero, but floating-point zero times infinity produces NaN. Selecting zero at an invalid position and multiplying an already computed quotient by zero are therefore not interchangeable numerical programs.

Inspect three objects separately: the raw dangerous-operation output, the values entering the final reduction, and the expressions needed for backward. A finite check on the scalar loss establishes only that this scalar is finite at that point. It does not replace intermediate-value checks or checks of the intended parameter gradients.

Do not define validity only as d!=0 when NaN is possible. NaN is unequal to zero and passes that simple predicate. The example requires finite numerators, finite denominators and nonzero denominators. Logarithms, square roots and other domains need their own conditions placed before the corresponding first unsafe operation.

Filter inputs first or protect operands first

One approach selects valid positions and divides only x[valid] by d[valid], then scatters into an initialized zero array if that layout is needed. Excluded positions do not undergo division. For the current fixed mask and fixed d, the objective is L=x₂/2 and its derivative with respect to x is (0,0.5), independently checkable using finite differences.

A second approach preserves a dense shape but prepares safe operands before division. Replace the numerator with zero and denominator with one at invalid positions, then divide. It produces the same (0,1.5) here. It neither clamps every real denominator to a positive epsilon nor changes the score at valid positions.

A nonzero but tiny denominator is a separate case. With x=3 and d=0.001, the quotient is 3000 and its derivative with respect to x is 1000; both are finite. Forcing the denominator to 0.01 changes those values to 300 and 100. The numbers become smaller, but the objective changes too. Finite does not mean appropriately scaled. A stabilization rule needs task-specific justification and should not be confused with excluding undefined positions.

Both methods require checks of actual shapes, broadcasting and types. Filtering can change the amount of data processed; dense protection retains safe work for excluded positions. Performance and memory depend on the complete path and target device. A short expression alone cannot establish a universal improvement.

Twelve checks distinguish arithmetic from autograd execution

The program actually runs NumPy arithmetic, explicitly computes the masked_chain formula, and checks both forms of pre-operation protection. The intentionally unsafe path locally suppresses warnings only for demonstration. The safe path instead raises on division or invalid-arithmetic errors. Silencing a warning is not the repair.

# NumPy arithmetic and explicit derivative formulas; no PyTorch autograd.
import numpy as np
x=np.array([1.,3.]); d=np.array([0.,2.])
valid=np.isfinite(x)&np.isfinite(d)&(d!=0)
with np.errstate(divide='ignore',invalid='ignore'):
    raw=x/d
    selected=np.where(valid,raw,0.)
    masked_chain=valid.astype(float)/d
    multiplied=raw*valid
assert np.isinf(raw[0])
assert selected.tolist()==[0.,1.5]
assert np.isfinite(selected.sum())
assert np.isnan(masked_chain[0]) and masked_chain[1]==.5
assert np.isnan(multiplied[0])

with np.errstate(divide='raise',invalid='raise'):
    safe=np.zeros_like(x)
    safe[valid]=x[valid]/d[valid]
    safe_x=np.where(valid,x,0.)
    safe_d=np.where(valid,d,1.)
    guarded=safe_x/safe_d
assert safe.tolist()==[0.,1.5]
assert np.array_equal(guarded,safe)
expected=np.zeros_like(x); expected[valid]=1/d[valid]
assert expected.tolist()==[0.,.5]

def fixed_mask_loss(z):
    return (z[valid]/d[valid]).sum()
for j in range(2):
    plus=x.copy(); minus=x.copy()
    plus[j]+=1e-5; minus[j]-=1e-5
    numerical=(fixed_mask_loss(plus)-fixed_mask_loss(minus))/(2e-5)
    assert np.isclose(numerical,expected[j],rtol=1e-8,atol=1e-8)
empty=np.zeros_like(valid,dtype=bool)
assert (x[empty]/d[empty]).sum()==0.
assert not (np.isfinite(np.array([np.nan])) & (np.array([np.nan])!=0)).any()
print('12 arithmetic and fixed-mask derivative checks passed; NumPy',np.__version__)

Both finite-difference checks perturb only x, keeping valid and d fixed. They test this step’s fixed-mask scalar objective. They do not differentiate the condition or establish smoothness as a denominator crosses zero. PyTorch is unavailable in the environment, so torch.where, automatic differentiation and GPU kernels were not executed.

Define what an all-masked training step means

When no valid entries remain, a sum can be defined as zero, but a mean over valid entries needs an explicit empty-set policy. Do not allow a zero count to enter the averaging denominator and hope that a final mask repairs it. Include an entirely excluded input as its own test rather than testing only mixtures of valid and invalid entries.

If the application chooses to skip the step, make the relevant update control flow explicit. A zero loss alone does not prove that parameters stay unchanged. Optimizer state and weight decay have separate semantics, illustrated by the earlier None versus zero-gradient update comparison. No single empty-batch policy is presented here as a universal training rule.

If the condition depends on model outputs, distinguish selecting the current piecewise path from proving differentiability at the selection boundary. Repairing the numerical domain does not remove every boundary. State whether finite differences compare within a fixed region or deliberately cross that boundary.

Find the earliest invalid operation

Start with deterministic inputs small enough to reproduce the discrepancy. Record the valid positions and the first nonfinite intermediate, then inspect loss and relevant parameter gradients. Comparisons should preserve the mathematical objective, valid set and reduction. A change that also drops examples, replaces denominators or alters averaging cannot be accepted as equivalent merely because NaN disappears.

If x itself was produced by an earlier undefined operation, replacing it with where at this later stage may still be too late. Move protection to the earliest domain boundary. Repairing an output is not an automatic cleanup of every upstream graph operation.

The issue is related to frozen features and gradient boundaries, but asks a different question. Here valid positions must retain their gradients. Wrapping the entire computation in no_grad to hide NaN also changes the intended training path.

If these before-and-after counterexamples help locate a finite-loss, invalid-gradient failure, voluntary support through the section below is welcome. The reusable result is a three-way comparison between an unsafe path, pre-operation protection and fixed-mask derivative checks.

Sources and validation scope

Autograd division-by-zero note · torch.where · Python expression evaluation · NumPy where

Checked against PyTorch 2.8 definitions and the autograd notes updated June 16, 2025. Diagrams and extended examples are original. NumPy arithmetic and explicit derivative checks ran; PyTorch did not. Planned article date: October 21, 2026 at 23:30 Beijing. Scheduling or actual public release is verified separately.