JAX PR #40814: Why select_n Returned a Value but Lost Its Gradient

JAX PR #40814 fixes a disagreement between selecting a value and sending its gradient back to the input that supplied it.

For an out-of-range integer selector, the reviewed forward implementation of lax.select_n chooses an endpoint case. Its old reverse-mode rule instead looked for an exact match between the original selector and each valid case number. No match meant zero gradients for every case, including the one returned by the forward calculation.

The patch clamps the selector inside the reverse rule. It also handles a single candidate directly, preventing a boolean selector from sending that rule into a nonexistent second case.

There is an important API boundary: JAX documents out-of-range select_n behavior as implementation-defined. This change makes reverse-mode differentiation agree with the reviewed lowering; it does not establish a new public guarantee that invalid indices will always be clipped.

A returned value with no apparent dependency

The mismatch can be reduced to two scalar inputs:

import jax
import jax.numpy as jnp

# Deliberately outside the valid two-case range, to inspect this implementation.
def choose(a, b):
    return jax.lax.select_n(jnp.int32(2), a, b)

a = jnp.float32(2.0)
b = jnp.float32(5.0)

print(choose(a, b))
print(jax.jacfwd(choose, argnums=(0, 1))(a, b))
print(jax.jacrev(choose, argnums=(0, 1))(a, b))

A local CPU check for this review used installed JAX and jaxlib 0.9.0.1, whose transpose rule still had the pre-fix logic. The equivalent scalar-loss probe returned:

value:                 5.0
forward-mode gradient: (0.0, 1.0)
reverse-mode gradient: (0.0, 0.0)

For this fixed selector, the implemented forward function returns b. Changing b changes the output by the same amount, so its derivative with respect to b is one. The integer selector itself is held constant; this is not an attempt to differentiate a discrete choice.

Issue #10993 reported a related discrepancy through lax.select, using the maximum uint32 value as its selector. That issue contains several other numerical examples. This patch addresses the selection rule, not every discrepancy listed there.

Forward lowering chooses by comparisons

The integer lowering in jax/_src/lax/lax.py constructs a binary comparison tree.

With two cases, its behavior is equivalent to:

which < 1 → case 0
otherwise → case 1

A negative signed index therefore reaches the first case. An index above the last valid case reaches the last case. With more cases, the tree recursively divides the candidates until one remains.

There is no explicit bounds-checking exception in this lowering. Endpoint behavior emerges from the comparisons and the single-case leaves.

Forward-mode automatic differentiation reuses select_n to choose among the input tangents. It therefore follows the same implemented routing as the primal value.

Reverse mode has a different task: distribute an incoming output derivative back to the appropriate input. That separate rule is where the two implementations diverged.

Reverse mode required an exact index match

JAX calls the incoming derivative a cotangent. For a scalar output used directly as a loss, that incoming cotangent is one.

The old integer transpose rule effectively computed this for every case i:

case i receives the cotangent when which == i
otherwise case i receives zero

For two cases and which=2, both comparisons are false:

2 == 0 → false
2 == 1 → false

The forward calculation had selected case 1, but the reverse calculation treated the selector as matching no case at all.

This explains the silent failure. Every returned gradient had a valid shape and floating-point type. Nothing in the result's structure revealed that the dependency on the selected input had disappeared.

The defect was in gradient routing, so increasing floating-point precision would not repair it.

The reverse rule now normalizes the selector first

The integer branch gains one line:

which = clamp(_const(which, 0), which, _const(which, len(cases) - 1))

The existing equality-based distribution then runs with an in-range selector.

For the two-case example:

original selector: 2
clamped selector:  1
cotangents:        (0, 1)

For a selector of -1, the result routes to case 0 instead.

_const constructs bounds using the selector's type. Signed and unsigned selectors therefore retain their own comparison domain; the fix does not first reinterpret a large unsigned value as a negative signed index.

No forward-lowering or JVP code changes in this PR. The new reverse rule is aligned with the behavior those paths already implement.

One candidate needs no two-way routing

The second defect occurs even with a single candidate and a boolean False selector, corresponding to valid case zero:

def only_case(x):
    return jax.lax.select_n(jnp.bool_(False), x)

The forward lowering returns the sole candidate directly.

The old boolean transpose branch assumed two cases and accessed cases[1]. Reverse differentiation could consequently raise IndexError even though evaluating the function succeeded.

The patch inserts an early return before that branch:

elif len(cases) == 1:
    return (None, ct if ad.is_undefined_primal(cases[0]) else None)

The leading None means there is no cotangent for the selector. When the case needs a cotangent, it receives the incoming one unchanged.

Here, is_undefined_primal is JAX's internal marker for an input being transposed through. It does not mean that the user's array contains undefined data.

The pre-existing symbolic-zero shortcut remains first, and the ordinary two-case boolean path remains unchanged.

The regression checks both differentiation modes

The added testSelectNGrad constructs scalar cases and arrays with shapes (3,) and (2, 3). Selectors are either scalars or match the case shape.

Its parameterization includes boolean, signed int32, and unsigned uint32 selectors; one and two candidates for all three; and six candidates for the integer types. The signed random range includes -1 and values beyond the final case. The unsigned range includes upper-bound violations.

It then calls:

check_grads(select_n, args, 2, ["fwd", "rev"], eps=1.)

The selector is captured by the test function. Differentiation and finite differences apply to the floating-point cases, with checks through second order.

Checking only the forward output would miss the zero-gradient defect. Checking only forward-mode derivatives would also miss it because the JVP already uses the same selection operation as the primal calculation.

This test exercises the relationship between those paths. It is still sampled coverage: JAX's check_grads uses a randomly chosen direction, not an exhaustive construction of every Jacobian entry. The PR does not attach a complete per-device execution report.

What was reproduced locally

The local investigation went beyond source inspection without claiming to build the merged tree.

First, the installed JAX 0.9.0.1 rule was exercised on CPU. Then only the select_n transpose-registry entry was temporarily replaced with the rule read from merge commit 8c7be284e2ec119d6126d41398147b6307e5d92b. Compilation caches were cleared between phases, and the original registry entry was restored afterward. No installed package files were edited.

Across 17 deterministic scenarios, the original rule produced nine gradient disagreements and two single-case IndexError failures. With the substituted rule, all 17 matched the forward routing, including JIT-compiled reverse gradients.

Representative scalar results were:

Selector and case count	Original reverse result	With reviewed rule
`-1`, two cases	`(0, 0)`	`(1, 0)`
`2`, two cases	`(0, 0)`	`(0, 1)`
maximum `uint32`, two cases	`(0, 0)`	`(0, 1)`
`False`, one case	`IndexError`	`(1,)`
`1`, two cases	`(0, 1)`	`(0, 1)`

A separate squared-selection probe produced the expected second derivative of two for the selected input under the substituted rule.

These are checks of the isolated rule running inside an older installed JAX runtime. They are not execution of the PR's original test suite or validation of the full merged commit on GPU or TPU.

Keep the application's index policy explicit

The reviewed rule now sends cotangents to the same input selected by the forward implementation, and the single-case path avoids indexing a second candidate.

Application code should still respect the documented selector range. When endpoint clipping is the intended behavior, make it explicit before calling select_n:

# cases must be nonempty, with matching shapes and dtypes.
which = jnp.clip(which, 0, len(cases) - 1)
result = jax.lax.select_n(which, *cases)

For lax.select, use a boolean predicate as documented.

The useful review principle is to check selection twice: which input supplied the value, and which input receives the derivative. In PR #40814, those answers finally agree at the reviewed boundaries.