ARCHIVED TRAINING / jax-ml/jax

jax-ml__jax-25957

This is a preserved training example. Membership in the current training wave has not been verified.

Problem statement

jax_debug_key_reuse incorrectly raises error
### Description

The following code raises error `KeyReuseError: Previously-consumed key passed to random_split at index 0` despite that we haven't used key2 yet.

```python
import jax; jax.config.update('jax_debug_key_reuse', True)
key = jax.random.key(0)
key1, key2 = jax.random.split(key)
_, key2 = jax.device_put((key1, key2))
jax.random.split(key2)
```

### System info (python version, jaxlib version, accelerator, etc.)

```
jax:    0.4.36
jaxlib: 0.4.36
numpy:  1.26.4
python: 3.12.7 (main, Oct  1 2024, 02:05:46) [Clang 16.0.0 (clang-1600.0.26.3)]
device info: cpu-1, 1 local devices"
process_count: 1
platform: uname_result(system='Darwin', node='phandu-mac.roam.internal', release='24.2.0', version='Darwin Kernel Version 24.2.0: Fri Dec  6 19:02:41 PST 2024; root:xnu-11215.61.5~2/RELEASE_ARM64_T6030', machine='arm64')
```

Training membership

opd_e4b_pure_v6a opd_e4b_pure_v6b opd_e4b_self_audit_v4

Task issues extracted from preserved verified_rft_messages training examples. This is not a verified current RL wave manifest.

Task and snapshot provenance

Snapshot captured 2026-11-12T17:06:59.476467+00:00 · Autoresearch-Arena/gemma-4-developer-agent

{
  "datasets": [
    "opd_e4b_pure_v6a",
    "opd_e4b_pure_v6b",
    "opd_e4b_self_audit_v4"
  ],
  "examples": 5,
  "instance_id": "jax-ml__jax-25957",
  "problem_statement": "jax_debug_key_reuse incorrectly raises error\n### Description\n\nThe following code raises error `KeyReuseError: Previously-consumed key passed to random_split at index 0` despite that we haven't used key2 yet.\n\n```python\nimport jax; jax.config.update('jax_debug_key_reuse', True)\nkey = jax.random.key(0)\nkey1, key2 = jax.random.split(key)\n_, key2 = jax.device_put((key1, key2))\njax.random.split(key2)\n```\n\n### System info (python version, jaxlib version, accelerator, etc.)\n\n```\njax:    0.4.36\njaxlib: 0.4.36\nnumpy:  1.26.4\npython: 3.12.7 (main, Oct  1 2024, 02:05:46) [Clang 16.0.0 (clang-1600.0.26.3)]\ndevice info: cpu-1, 1 local devices\"\nprocess_count: 1\nplatform: uname_result(system='Darwin', node='phandu-mac.roam.internal', release='24.2.0', version='Darwin Kernel Version 24.2.0: Fri Dec  6 19:02:41 PST 2024; root:xnu-11215.61.5~2/RELEASE_ARM64_T6030', machine='arm64')\n```",
  "repo": "jax-ml/jax",
  "source": "v2/experiments/train_data/opd_e4b_pure_v6a/verified_rft_messages.jsonl",
  "source_sha256": "f045b7ca5b3883a1e2eaee81e1c02a4a1b53dd7b7b36e98bacb2ba5823a15148"
}

{
  "captured_at": "2026-11-12T17:06:59.476467+00:00",
  "membership": "archived training",
  "membership_evidence": "Task issues extracted from preserved verified_rft_messages training examples. This is not a verified current RL wave manifest.",
  "name": "Preserved post-training task archive",
  "source": "Autoresearch-Arena/gemma-4-developer-agent",
  "source_revision": "016496de222aaafca0d4ee4c18efc63ba616f76e",
  "sources": [
    {
      "path": "v2/experiments/train_data/opd_e4b_cispo_v5/verified_rft_messages.jsonl",
      "sha256": "a43610ad0dfb6ace95375cfae09000f41afc0338cc925af870a4ade7a9b20c4c"
    },
    {
      "path": "v2/experiments/train_data/opd_e4b_hybrid_v3/verified_rft_messages.jsonl",
      "sha256": "bd326303eb127e8e5e08f03daf97bb88d79d19709d59f3f93f348bc56686a8df"
    },
    {
      "path": "v2/experiments/train_data/opd_e4b_onpolicy_v3/verified_rft_messages.jsonl",
      "sha256": "960c680e797e823a0d3a22f4b1d5ef60a72971373d703b877340b0c9126aaf68"
    },
    {
      "path": "v2/experiments/train_data/opd_e4b_pure_v6a/verified_rft_messages.jsonl",
      "sha256": "f045b7ca5b3883a1e2eaee81e1c02a4a1b53dd7b7b36e98bacb2ba5823a15148"
    },
    {
      "path": "v2/experiments/train_data/opd_e4b_pure_v6b/verified_rft_messages.jsonl",
      "sha256": "0f62460a9706eea5e1605255896f191ab00cb5706c5f1f2e1c293e9a718218dc"
    },
    {
      "path": "v2/experiments/train_data/opd_e4b_pure_v7/verified_rft_messages.jsonl",
      "sha256": "e37f7bd675322f3f81500b73559d60d5f602ca0537bf19f5cb06332c728c7b72"
    },
    {
      "path": "v2/experiments/train_data/opd_e4b_pure_v7b/verified_rft_messages.jsonl",
      "sha256": "b864bdc8695183433eef6e31cb8d33fad44ba3aded2170c1ade4ad5b2944ce23"
    },
    {
      "path": "v2/experiments/train_data/opd_e4b_self_audit_v4/verified_rft_messages.jsonl",
      "sha256": "0026023745afb8f3dc1c176dc7f7171157b4ad5a66656b524af218e8e19ceee6"
    },
    {
      "path": "v2/experiments/train_data/rft_e4b_aligned_v2/verified_rft_messages.jsonl",
      "sha256": "92a4515fee9d152b0d4c510ee3ce4caaaffe2003ba7584469f8027f174aab18c"
    },
    {
      "path": "v2/experiments/train_data/rft_v1/verified_rft_messages.jsonl",
      "sha256": "a42453ccd7ac7680b0a03865e8e7080bc8bf74c8d9d193b440c441ed46f8512b"
    },
    {
      "path": "v2/experiments/train_data/rft_v2/verified_rft_messages.jsonl",
      "sha256": "18f4b3bd7e651f83c6acab479c9475e3fa50dc652a384f309e97146983ac85ba"
    }
  ]
}