Fix eager error_if hanging when inputs are sharded over multiple devices

[discobot]

Fixes #1232.

The hang happens because eager error_if routes through filter_jit, so a
multi-device-sharded input compiles an SPMD program over every device in the
sharding. The all-reduce in the abort trace sits inside the lax.cond error
branch: once the exception from jax.pure_callback kills one device's
execution thread, its peers never reach the collective and XLA aborts the
process when the rendezvous times out.

This PR handles that case before the jit: if every input is concrete (no
tracers) and fully addressable, and at least one is sharded over more than one
device, then the predicate is evaluated on the host. On error, the inputs are
copied to the host and passed to the existing branched_error_if_impl_jit,
which then compiles a single-device program -- so the usual
EquinoxRuntimeError reporting is unchanged. For the non-raising on_error
modes, each output leaf's original sharding is restored with jax.device_put.
Traced or non-fully-addressable inputs fall through to the current behaviour.

One known limitation: error_if under a user's jax.jit over multi-device
inputs still aborts in the same way (verified locally). There the multi-device
program is what the user asked for, and the underlying problem -- a raising
callback inside an SPMD program -- doesn't look fixable from inside error_if.
This PR covers eager mode, as reported.

Tested with the two CPU devices already configured in tests/conftest.py: the
new test aborts the process (exit 134) without the fix and passes with it,
covering replicated and non-replicated shardings, each of x/pred sharded
alone, and sharding preservation on the false-predicate and
on_error="nan"/"off" paths.

[patrick-kidger]

Hi there! Thank you for this, and sorry for taking a few weeks to get around to it :)

The reason for the delay is that we've been beavering away at landing a number of other improvements to error_if – see #762, just merged into dev – and I wanted to land that first :)

Would you be interested in rebasing this branch on top of dev, and seeing how things go?

I need to stare at the implementation a little bit – I'm not an expert on multi-device JAX – but I don't love the special-casing of this case. Would it be possible to just 'always' move to a single device... which will be a no-op in the inherently-single-device case?

[discobot]

Thanks, and no worries about the delay! Rebased onto dev and reworked it the way you suggested. No more special-casing of the multi-device path: _error_impl now always moves the eager check onto a single device (a no-op when the inputs already live on one) and then restores each output's original sharding. error_if, warn_if and branched_error_if all route through _error_impl, so all three are covered, and the regression test uses the two CPU devices from conftest.

One judgement call worth your eye: I restore the input sharding onto the output, so error_if stays transparent in the non-raising nan/off modes. Happy to drop that if you'd rather the multi-device case just return single-device.