Sixteen JAX devices on a CPU, and one gradient bug worth finding

My favorite use for a pretend cluster is a real correctness failure. Here, replacing an average with...