RuntimeError when running tutorial notebook on Distributed Inference with JAX

I'm trying to run the Distributed Inference with JAX tutorial notebook on a p3.16xlarge ec2 instance (with 8 GPUs) running the Deep Learning AMI.
After successfully installing TF, TFP and Jax on a fresh environment, I am able to run some of the first cells, like

# check if Jax detectects the GPUs
print(f'Found {jax.device_count()} devices')
# Found 8 devices

xs = jnp.arange(8.)
out = jax.pmap(lambda x: x + 5.)(xs)
print(type(out), out)
# <class 'jax.interpreters.pxla.ShardedDeviceArray'> [ 5.  6.  7.  8.  9. 10. 11. 12.]

but when I try to run the following cell I get an error

def f(x):
    out = lax.psum(x, axis_name='i')
    return out
xs = jnp.arange(8.) # Length of array matches number of devices
jax.pmap(f, axis_name='i')(xs)

The error I get is

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-9-dabe8144040f> in <module>()
      3     return out
      4 xs = jnp.arange(8.) # Length of array matches number of devices
----> 5 jax.pmap(f, axis_name='i')(xs)

/home/ubuntu/anaconda3/envs/jax/lib/python3.6/site-packages/jax/api.py in f_pmapped(*args, **kwargs)
   1569         out_axes_thunk=out_axes_thunk,
   1570         name=flat_fun.__name__, donated_invars=tuple(donated_invars),
-> 1571         global_arg_shapes=tuple(global_arg_shapes_flat))
   1572     return tree_unflatten(out_tree(), out)
   1573 

/home/ubuntu/anaconda3/envs/jax/lib/python3.6/site-packages/jax/core.py in bind(self, fun, *args, **params)
   1459   def bind(self, fun, *args, **params):
   1460     assert len(params['in_axes']) == len(args)
-> 1461     return call_bind(self, fun, *args, **params)
   1462 
   1463   def process(self, trace, fun, tracers, params):

/home/ubuntu/anaconda3/envs/jax/lib/python3.6/site-packages/jax/core.py in call_bind(primitive, fun, *args, **params)
   1391   tracers = map(top_trace.full_raise, args)
   1392   with maybe_new_sublevel(top_trace):
-> 1393     outs = primitive.process(top_trace, fun, tracers, params)
   1394   return map(full_lower, apply_todos(env_trace_todo(), outs))
   1395 

/home/ubuntu/anaconda3/envs/jax/lib/python3.6/site-packages/jax/core.py in process(self, trace, fun, tracers, params)
   1462 
   1463   def process(self, trace, fun, tracers, params):
-> 1464     return trace.process_map(self, fun, tracers, params)
   1465 
   1466   def post_process(self, trace, out_tracers, params):

/home/ubuntu/anaconda3/envs/jax/lib/python3.6/site-packages/jax/core.py in process_call(self, primitive, f, tracers, params)
    598 
    599   def process_call(self, primitive, f, tracers, params):
--> 600     return primitive.impl(f, *tracers, **params)
    601   process_map = process_call
    602 

/home/ubuntu/anaconda3/envs/jax/lib/python3.6/site-packages/jax/interpreters/pxla.py in xla_pmap_impl(fun, backend, axis_name, axis_size, global_axis_size, devices, name, in_axes, out_axes_thunk, donated_invars, global_arg_shapes, *args)
    617                                    donated_invars, global_arg_shapes,
    618                                    *abstract_args)
--> 619   return compiled_fun(*args)
    620 
    621 @lu.cache

/home/ubuntu/anaconda3/envs/jax/lib/python3.6/site-packages/jax/interpreters/pxla.py in execute_replicated(compiled, backend, in_handler, out_handler, *args)
   1128 def execute_replicated(compiled, backend, in_handler, out_handler, *args):
   1129   input_bufs = in_handler(args)
-> 1130   out_bufs = compiled.execute_sharded_on_local_devices(input_bufs)
   1131   if xla.needs_check_special():
   1132     for bufs in out_bufs:

RuntimeError: Internal: external/org_tensorflow/tensorflow/compiler/xla/service/gpu/nccl_utils.cc:193: NCCL operation ncclCommInitRank(&raw_comms[i], num_participants, unique_id, local_participants[i].rank) failed: unhandled cuda error: while running replica 0 and partition 0 of a replicated computation (other replicas may have failed as well).

This is how I set up the environment and install the packages.

# create fresh environment
conda create --name jax --clone python3
conda activate jax

# switch to from cuda 10.0 to CUDA 11.0 as per the DLAMI documentation
sudo rm /usr/local/cuda
sudo ln -s /usr/local/cuda-11.0 /usr/local/cuda

# install packages
pip install --upgrade pip
pip install jaxlib --upgrade 
pip install tfp-nightly[jax] --upgrade
pip install tf-nightly
pip install --upgrade jax jaxlib==0.1.65+cuda110 -f https://storage.googleapis.com/jax-releases/jax_releases.html

The failing entry point is the Distributed Inference with JAX tutorial notebook, specifically the jax.pmap cell calling lax.psum. Start by reproducing that cell on the stated 8-GPU setup and checking the listed JAX, JAXLIB, TensorFlow, TFP, and CUDA installation steps. Done means the cell runs successfully or the notebook documents the supported environment and limitation.