Pass in-memory dataset as a dynamic JAX sampling argument

[sylvesterkaczmarek]

Fixes #285.

Changes

• Stop closing over the entire in-memory dataset inside the JAX-jitted sampling function on the single-device path.
• Compile sample_per_shard(data, key) and supply the device-resident dataset as a dynamic argument at each call. This avoids embedding the dataset as an XLA constant while retaining existing batching, RNG evolution, and sampling behavior.
• Preserve the separate multi-device pmap path unchanged.
• Add a regression test that verifies the JIT function accepts independent data and key arguments, then checks a real CPU JAX sample.

Validation

• The actual modified iterator class and new regression test were run using CPU JAX 0.10.2 and small dependency stubs: 1 regression test passed. Two additional consecutive samples had the expected (2, 3) shape with dataset size 8.
• python -m compileall on changed files: passed.
• ruff check --select E4,E7,E9,F on changed files: passed.
• git diff --check: passed.

The full Acme TensorFlow/TFDS integration test suite and GPU memory behavior were not exercised locally; maintainers should confirm them in CI or a suitable runtime.