Distributed-scale solve fixes: sharded NaN/Inf reductions

[hmgaudecker]

Re-implements the still-relevant distributed-scale solve fixes from the abandoned #364, fresh on top of the discrete-only-sharding work (#368, now merged to main). Verified on 4 CPU devices (jax_num_cpu_devices=4).

What changed

Sharded NaN/Inf reductions (code). validate_V, contains_nan, log_nan_in_V, and solve's hot-loop diagnostic accumulator route their NaN/Inf checks through module-level @jax.jit helpers v_array_has_nan / v_array_has_inf. GSPMD then partitions each reduction across the V-array's devices (per-device any → all-reduce → replicated scalar). The eager jnp.any(jnp.isnan(V_arr)) alternative can fall back to gathering a sharded V onto the default device before reducing — exhausting device memory at production grid sizes. Regression: two checks that both helpers return a fully-replicated 0-d scalar on a NamedSharding-sharded input.

Batching/sharding tuning guide (docs). New docs/user_guide/tuning.md section contrasting the two grid knobs — batch_size (splay: a time-neutral memory knob) and distributed (shard: the only wall-clock knob, discrete non-transitioning axes only) — when each applies, and the boundaries the constructor enforces (continuous axes can't shard; you can't splay and shard the same axis). The reasoning behind the subtlest part of that guide is written out below — it documents existing behaviour, there is no code change for it.

Rationale for the guide — batch_size chunks the state sweep, not the shock expectation

(Background, not a code change: this is the reasoning the new guide distills.)

Each period's backward step is two nested sweeps:

1. State sweep. Walk the grid of current states; for each current state compute V = max_a Q. A grid's batch_size chunks this sweep. The chunked axis is a current state — a dimension of the output V — so each chunk's results are written out and freed, and peak memory shrinks.
2. Shock expectation. Inside one Q(state, action) cell, average the next period's value over the stochastic transition: E[V'] = Σ_{s'} P(s')·V'(s', …), built as a productmap over the next-period stochastic outcomes followed by jnp.average.

A stochastic grid (e.g. health) names an axis in both: a current-period state in the state sweep, and a summed-over next-period outcome in the shock expectation. batch_size acts only on the state sweep — tests/test_stochastic.py::test_stochastic_state_batch_size_is_value_equivalent_to_no_splay pins that solving with batch_size=1 gives a value function identical to batch_size=0.

The shock expectation stays unconditionally fused (batch_size=0). Its axis is summed away, not kept, so chunking it can't shrink the output — it only de-fuses the jnp.average from the V' interpolation, forcing the full stack of per-outcome V' arrays to materialize before the sum instead of streaming into the running total. Same jax.lax.map(batch_size=1) primitive, opposite sign depending on whether the chunked axis is kept (an output) or summed away (reduced) — isolated lax.map + jnp.average probe, K=8 next-period outcomes over an N=4M-cell working set (128 MB = K·N·4B):

batch_size	summed-away axis (shock expectation)	output axis (state sweep)
0 (→ vmap)	0 MB — fused into the average, streamed	128 MB — all transients live
1 (→ scan)	128 MB — full [K, N] stack materialized before the sum	16 MB — one transient at a time

Scale N to the production state-action working set and K to the joint next-period stochastic cardinality and 0 → 128 MB becomes fits → 271 GiB. To shed memory along a stochastic dimension, shard the current axis (distributed=True) or chunk the state sweep — both keep the shock expectation fused and beat a sequential scan.

Dropped from #364

out_shardings on the solve lowering (#364 piece 1) — obsolete. Its motivating case (continuous distributed=True) is forbidden outright (GridInitializationError), and discrete-distributed V already shards correctly via natural jit output inference (verified: num_devices=4 with and without out_shardings; all distributed tests pass without it). A docstring note on _fail_if_continuous_grid_distributed records the prerequisite (and points at #364's reference implementation) should continuous-axis sharding ever be re-enabled.

Device-count padding (#364 piece 6) landed in #368.

Verification

• tests/test_stochastic.py (incl. the new batch_size-equivalence test) — relevant tests pass locally
• pixi run -e type-checking ty — clean
• prek run --all-files — clean
• Full suite (pytest tests -n 4; tests/test_distributed.py on 4 CPU devices) — deferred to CI

🤖 Generated with Claude Code

[mj023]

Good change, very helpful guide about batching and sharding!

The second half of the PR description reads like there is a second change, but I can't find it in the Code?

[mj023]

I cant quite follow what this means. To shard a second axis, one would need to add devices anyways, rigth?

[hmgaudecker]

Thanks, that is misleading indeed. Now that we disallow sharding continuous axes, you may well want to shard 2x2, so that's supported. But in that example, it did not make any sense.

[hmgaudecker]

The second half of the PR description reads like there is a second change, but I can't find it in the Code?

That was the rationale more than anything, poor cleaning up of the writing from here. Should be better now.