vmap HkDisentangledRNN.update_latents per-latent loop (OOM / PJRT arg-packing ceiling at latent_size=256)

## Context

Study 10 (`10-disrnn-near-gru-ceiling`, stage3-wide-latent) tried to launch a disRNN
with `latent_size=256` (matching study 01's H=256 GRU capacity exactly) on a single
H200 (141GB). Both tasks of the 2-seed grid (Beaker exp
[01M22Q2AF2JNNJP69RWQ8CG4AA](https://beaker.org/ex/01M22Q2AF2JNNJP69RWQ8CG4AA))
failed after ~53 min at JIT-compile time with:

```
RESOURCE_EXHAUSTED: Out of memory while trying to allocate 1513245575640 bytes
```

(~1.38 TiB requested vs 141GB available). `hlo_rematerialization.cc` logged that its
own rematerialization pass could only reduce the graph from 1.38TiB to 1.38TiB — i.e.
the graph itself, not just activations, is enormous.

A follow-up diagnostic (batch_size 1024 -> 64, same `latent_size=256`, one-off exp
[01M23R8FKZSW7D1CDQZFSNDGW4](https://beaker.org/ex/01M23R8FKZSW7D1CDQZFSNDGW4)) ruled out
activation memory as the dominant term: it did not reproduce the OOM, but hit a
*different*, batch-size-independent failure at the same `train_step` JIT-compile point:

```
INVALID_ARGUMENT: Can't pack device memory arguments array of size 1793 which is
larger than the maximum supported size of 1024
```

## Root cause

`HkDisentangledRNN.update_latents` (`disentangled_rnns/library/disrnn.py`) builds
**one separate `ResMLP` per latent** in a Python-level `for net_i in
range(latent_size)` loop — unrolled at JAX trace time, not `vmap`ped. At
`latent_size=256` (vs. `latent_size=5` used by every prior study) this is a ~51x
increase in the number of distinct update-net modules traced into one XLA graph, each
contributing its own parameter/gradient/optimizer-state buffers as separate packed
device-memory arguments to the compiled `train_step`. That blows both:

1. the GPU memory budget (huge unrolled graph, rematerialization barely helps), and
2. a fixed PJRT argument-packing ceiling (1024 packed device-memory arguments per
compiled call) — independent of batch size, so no batch-size reduction can fix it.

## Findings

| Config | batch_size | Result |
|---|---|---|
| `latent_size=256` | 1024 | `RESOURCE_EXHAUSTED`, ~1.38TiB requested (exp 01M22Q2AF2JNNJP69RWQ8CG4AA) |
| `latent_size=256` | 64 | `INVALID_ARGUMENT`, packed-args=1793 > 1024 limit (exp 01M23R8FKZSW7D1CDQZFSNDGW4) |
| `latent_size=32` | 1024 | Cleared JIT compile + first training step (exp 01M23VHJSCFKF6Y80WDZGSJCZ0) |

Study 10's stage3-wide-latent was relaunched at the downgraded `latent_size=32`
(Beaker exp 01M23ZX664WP1QJQJNVBYY1V13) as a workaround; this does **not** reach the
study's original GRU-H256-matched capacity target.

## Fix needed

`vmap` the per-latent loop in `HkDisentangledRNN.update_latents` instead of
Python-unrolling it, so the traced graph size (and packed-argument count) is
independent of `latent_size`. This should let `latent_size=256` (and beyond) compile
and fit in GPU memory the same way `latent_size=5` does today.

## Done when

- [ ] `update_latents`'s per-latent `ResMLP` application uses `jax.vmap` (or an
equivalent batched formulation) instead of a Python `for` loop over latents.
- [ ] A disRNN with `latent_size=256` (study 10's original stage3-wide-latent config)
trains a full step without `RESOURCE_EXHAUSTED` or PJRT argument-packing errors
on a single H200 (141GB), at the study's intended `batch_size=1024`.
- [ ] Existing disRNN behavior at `latent_size=5` is unchanged (numerically
equivalent likelihood / bottleneck metrics) after the vmap refactor — a
regression test comparing pre/post-fix outputs on a small fixed seed is
sufficient.
- [ ] `TRAINING.md` / inline comments in `disrnn.py` note the vmap change and its
motivation (this issue).

## Notes

- Referenced study: `aind-dynamic-foraging-bfm-dispatcher`
`studies/10-disrnn-near-gru-ceiling/variants/stage3-wide-latent/` (branch
`study/10-disrnn-near-gru-ceiling`).
- Intervention record with full evidence:
`studies/10-disrnn-near-gru-ceiling/variants/stage3-wide-latent/launch_record/beaker_resubmit_latent32.json`.
- Verified (not suspected): both the OOM and the argument-packing failure were
observed directly in Beaker job logs (see linked experiment ids above).