`RandomState` can not be treated as a dynamical variable in parallelized JIT compilation phase.

brainpy / BrainPy

Brain Dynamics Programming in Python

https://brainpy.readthedocs.io/

GNU General Public License v3.0

531 stars 94 forks source link

`RandomState` can not be treated as a dynamical variable in parallelized JIT compilation phase. #469

Closed chaoming0625 closed 10 months ago

chaoming0625 commented 1 year ago

Please:

[x] Check for duplicate requests.
[x] Describe your goal, and if possible provide a code snippet with a motivating example.

Currently, the random key in a RandomState is not well compatible with the pjit method, because its sharding strategy is not defined. Also, random key is not compatible with parallelization APIs like vmap and pmap.

One promising strategy is to specially treat RandomState: only pass the needed key into the function without new key returns.

chaoming0625 commented 10 months ago

Has been supported.