[JAX] Fix undefined sr_rng_state that breaks --dry-run in two encoder examples

[Anai-Guo]

Description

examples/jax/encoder/test_model_parallel_encoder.py and
test_multiprocessing_encoder.py build the RNG dict for the --dry-run path
from a name that is never bound:

if args.dry_run:
    labels = jnp.zeros(label_shape, dtype=jnp.bfloat16)
    rngs = {DROPOUT_KEY: dropout_rng, SR_KEY: sr_rng_state}   # sr_rng_state is undefined
    jit_train_step(state, inputs, masks, labels, var_collect, rngs)
    print("PASSED")
    return None

The local inside train_and_evaluate is sr_rng (bound a few dozen lines
earlier alongside params_rng / dropout_rng), and sr_rng_state is a name
from the quantizer API (transformer_engine/jax/cpp_extensions/quantization.py),
not from these scripts. So

python examples/jax/encoder/test_model_parallel_encoder.py --dry-run
python examples/jax/encoder/test_multiprocessing_encoder.py --dry-run

both die with NameError: name 'sr_rng_state' is not defined before the single
train step ever runs.

$ pyflakes examples/jax/encoder/test_model_parallel_encoder.py
...:385:47: undefined name 'sr_rng_state'
$ pyflakes examples/jax/encoder/test_multiprocessing_encoder.py
...:482:47: undefined name 'sr_rng_state'

The other two examples already do it correctly

The same block in the two sibling examples uses sr_rng:

file	line	code
test_single_gpu_encoder.py	254	rngs = {DROPOUT_KEY: dropout_rng, SR_KEY: sr_rng}
test_multigpu_encoder.py	352	rngs = {DROPOUT_KEY: dropout_rng, SR_KEY: sr_rng}
test_model_parallel_encoder.py	385	... SR_KEY: sr_rng_state} ←
test_multiprocessing_encoder.py	482	... SR_KEY: sr_rng_state} ←

The non-dry-run loop in the two broken files also uses sr_rng
(lines 395 and 491), so sr_rng is unambiguously the intended name.

Fix

sr_rng_state → sr_rng in the two --dry-run branches. Two lines.

---

🤖 Generated with Claude Code

[greptile-apps]

Confidence Score: 5/5

The PR appears safe to merge and directly fixes the dry-run failures without changing normal training behavior.

Summary

• Prevents NameError in the model-parallel encoder dry run.
• Prevents NameError in the multiprocessing encoder dry run.
• Aligns both branches with their normal training paths and sibling examples.