dkn16

## Code ``` from jax.config import config config.update("jax_enable_x64", True) import os #os.environ['XLA_PYTHON_CLIENT_ALLOCATOR']='platform' import jax import jax.numpy as jnp from tensorflow_probability.substrates import jax as tfp tfd = tfp.distributions Y=jnp.ones((4032,258),dtype=jnp.float64) distribution =...

Error when running shared HMC on distributed GPUs with big dataset

5

I'm currently running a distributed HMC on 4 Tesla V100 cards, and my codes are like: `import functools` `import collections` `import contextlib` `from jax.config import config` `config.update("jax_enable_x64", True)` `import jax`...