VonMisesFisher returns NaN with mean_direction=[1., 0., 0.] with lax backend

As in the title: VMF samples nan with mean_direction=[1., 0., 0.] and works with other unit directions. Example:

import tensorflow_probability.substrates.jax as tfp
from jax import random

vmf_good = tfp.distributions.VonMisesFisher(
    mean_direction=[0., 1., 0.],
    concentration=1.,
)

vmf_bad = tfp.distributions.VonMisesFisher(
    mean_direction=[1., 0., 0.],
    concentration=1.,
)

print(vmf_good.sample(sample_shape=2, seed=random.PRNGKey(0)))
print(vmf_bad.sample(sample_shape=2, seed=random.PRNGKey(0)))

>>>2021-01-04 18:39:35.373535: W tensorflow/stream_executor/platform/default/dso_loader.cc:60] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory
/home/jatentaki/Storage/jatentaki/miniconda3/envs/tfp/lib/python3.8/site-packages/jax/lib/xla_bridge.py:130: UserWarning: No GPU/TPU found, falling back to CPU.
  warnings.warn('No GPU/TPU found, falling back to CPU.')
[[-0.9372401  -0.17581224 -0.30111626]
 [-0.93092006 -0.36028326 -0.05986603]]
[[nan nan nan]
 [nan nan nan]]

originally reported in https://github.com/pyro-ppl/numpyro/issues/859

Run the provided Python reproduction with the JAX lax backend, then trace the VonMisesFisher sampling path for mean_direction=[1., 0., 0.]. Compare it with the working [0., 1., 0.] case and add a regression test showing that samples no longer contain NaN values for the axis-aligned direction.