Batch the backward pass instead of vmapping it

[yoch]

backward() is an unbatched function under jax.vmap, and its while_loop carries the whole tree. Under vmap the loop predicate becomes batched, so JAX turns every iteration into a select between the old and the new carry for every array of the tree. A backup step then costs O(num_simulations) instead of O(path length), and a whole search O(num_simulations²).

This version processes the batch in one loop with a scalar predicate (any(index != ROOT)). Rows that have reached the root rewrite their current values and are otherwise left alone; the other rows scatter the entries on their path. The trees are identical to the current ones (the new tests compare against the previous implementation).

Measured on CPU (4 cores), gumbel_muzero_policy with a trivial recurrent function, best of 5 calls:

batch, actions, simulations	main	this PR
16, 81, 32	6.6 ms	4.7 ms (1.4x)
16, 81, 128	90 ms	34 ms (2.7x)
16, 81, 400	544 ms	100 ms (5.5x)
8, 1729, 64	132 ms	50 ms (2.6x)
1, 81, 128	2.8 ms	2.9 ms (no change)

On a Colab T4 (jax 0.11.1), median of 5 calls, same benchmark:

batch, actions, simulations	main	this PR	#116
16, 81, 32	11.3 ms	10.2 ms (1.1x)	12.7 ms (0.9x)
16, 81, 128	50.6 ms	43.2 ms (1.2x)	45.4 ms (1.1x)
128, 1729, 32	83 ms	18.8 ms (4.4x)	20.9 ms (4.0x)
128, 1729, 64	295 ms	39.3 ms (7.5x)	44.8 ms (6.6x)

The gain grows with the number of actions, which is what the diagnosis predicts (the selects are over arrays of size [num_simulations, num_actions]); with 81 actions it is within the noise of a shared runtime. In every configuration, on CPU, GPU and TPU, the trees are identical to those of main.

On a Colab TPU v5e (jax 0.7.2, the oldest version pyproject.toml allows), median of 5 calls:

batch, actions, simulations	main	this PR	#116
16, 81, 32	2.44 ms	2.40 ms (1.02x)	7.38 ms (0.33x)
16, 81, 128	9.0 ms	8.9 ms (1.01x)	17.5 ms (0.51x)
128, 1729, 32	33.8 ms	30.5 ms (1.11x)	48.5 ms (0.70x)
128, 1729, 64	150 ms	103 ms (1.46x)	195 ms (0.77x)

So on TPU this is never slower than main and gains from 1.1x to 1.5x with many actions. #116 was slower than main in all four of my runs there, which matches the caution in its description about TPU; I only ran it as it stands, without tuning. A real network will dilute all of these numbers, since the cost of the model call is not in this benchmark.

Related: this touches the same function as #116, which attacks the same cost with a different structure (it keeps the vmap and records the path in a small carry, then does one scatter per array after the loop). On the same benchmarks the two are close for small trees and small action sets, and this one is ahead where #116's own description measures (batch 128, 1729 actions: 7.5x against 6.6x at 64 simulations on the T4; on CPU the difference grows with the number of simulations, 5.5x against 3.6x at 400), and on TPU only this version avoids a slowdown (table above). The two PRs conflict textually, and so do #128 and #129 (they edit the same region). I am happy to rebase on whichever lands first, or to help combine them if you prefer.

Relevant to #107 (loop_fn in backward slow as the tree gets deeper) only partly: this reduces the cost of each iteration, it does not reduce the number of iterations or kernel launches, which is what matters most on GPU. I would not close that issue with it.

Tests: pytest from _testing/: 29 passed (24 existing plus 5 new), pylint clean on the changed files.