You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
Currently the code runs on one device, which doesn't allow scaling to larger computational network such as TPU pods.
Parallelizing over local sampler should be relatively simple, since that does not required communication between devices. Note that if single evaluation of the likelihood demands more RAM than what's available on the chips (TPUv4 has 8GB RAM per core, gradient of functions may cause problem), the computation may need to be shard to multiple device, but that should be taken care separately.
Evaluation of global sampler should be similar to local sampler.
Currently the code runs on one device, which doesn't allow scaling to larger computational network such as TPU pods.
Parallelizing over local sampler should be relatively simple, since that does not required communication between devices. Note that if single evaluation of the likelihood demands more RAM than what's available on the chips (TPUv4 has 8GB RAM per core, gradient of functions may cause problem), the computation may need to be shard to multiple device, but that should be taken care separately.
Evaluation of global sampler should be similar to local sampler.
Training the normalizing flow requires collecting data from multiple devices and updating weights in a somewhat sync version. Have a look of
pmap
to see how to deal with that https://jax.readthedocs.io/en/latest/jax-101/06-parallelism.htmlThe text was updated successfully, but these errors were encountered: