Guard duplicate-key check and write in safetensors loader with a lock

[AMR5210]

Resolves #1259

This fixes a TOCTOU race in load_and_create_model_orig
(tunix/models/safetensors_loader.py). process_key runs across a
ThreadPoolExecutor(max_workers=os.cpu_count()), but the duplicate-key check
if jax_key_mapped in file_loaded_tensors and the following write were not
protected by a lock. Two threads whose source keys map to the same jax key can
both pass the check before either writes, so the second silently overwrites the
first and the intended ValueError never fires. The model then loads with
incorrect weights and no error.

The fix adds a dict_lock and wraps only the check and write in it. The tensor
read, transform, and dtype cast stay outside the lock, so loading remains
parallel with no meaningful speed impact.

Reference

• Issue: TOCTOU race condition in concurrent safetensors loading bypasses duplicate key detection #1259

Colab Notebook
N/A. This is a bug fix to existing loading code and adds no new API.

Checklist

• I have added all the necessary unit tests for my change.
• I have verified that my change does not break existing code and all unit tests pass.
• I have added all appropriate doc-strings/documentation.
• My PR is based on the latest changes of the main branch (if unsure, rebase the code).
• I have signed the Contributor License Agreement.
• I have followed Contribution Guidelines.

[AMR5210]

Opened #1627 to fix this. The fix adds a dict_lock around the duplicate-key
check and write in load_and_create_model_orig so both operations are atomic
under the existing ThreadPoolExecutor. Includes a regression test that
confirms the duplicate is detected deterministically. All 7 tests pass on
Kaggle TPU infrastructure.