Fix KeyReuseError

✅ Solution

The "KeyReuseError" in jax arises when the same PRNGKey is used multiple times within a jax transformation (like vmap or pmap) without being split, which leads to non-deterministic or potentially correlated random numbers. To fix this, ensure that each PRNGKey used within a transformed function is uniquely split using `jax.random.split` before each use, generating new keys for subsequent operations, effectively preventing key reuse. Always split your PRNGKey before passing it to JAX functions that use randomness within loops or vmapped code.