Repository: pyro-ppl/numpyro

Issue:
Support float64 array
Currently, we only develop and test for float32 array. When the functionalities are ready, we should test for float64 array too.


Discussion:
> When the functionalities are ready, we should test for float64 array too.

Note that float64 will be significantly slower than float32 for GPU, so I think it is a good practice to write our code assuming that we will run it on float32, i.e. ensure numerical stability whenever possible. e.g. by providing logit support in distributions. I think tests passing with float32 should almost always pass with float64 (the reverse isn't true)?
@neerajprad This is to prevent something like `a = np.array([1.])`, `np.arange(...)`, or `random.normal(...)` without `dtype` in the code.

Resolving patch:
diff --git a/numpyro/distributions/util.py b/numpyro/distributions/util.py
--- a/numpyro/distributions/util.py
+++ b/numpyro/distributions/util.py
@@ -8,6 +8,7 @@
 import jax.numpy as np
 from jax import canonicalize_dtype, custom_transforms, device_get, jit, lax, random, vmap
 from jax.interpreters import ad, batching
+from jax.lib import xla_bridge
 from jax.numpy.lax_numpy import _promote_args_like
 from jax.scipy.special import gammaln
 from jax.util import partial
@@ -183,7 +184,7 @@ def _standard_gamma_p(key, alpha):
 
 
 @partial(jit, static_argnums=(2, 3))
-def _standard_gamma(key, alpha, shape=(), dtype=np.float32):
+def _standard_gamma(key, alpha, shape, dtype):
     shape = shape or np.shape(alpha)
     alpha = lax.convert_element_type(alpha, dtype)
     if np.shape(alpha) != shape:
@@ -191,14 +192,15 @@ def _standard_gamma(key, alpha, shape=(), dtype=np.float32):
     return _standard_gamma_p(key, alpha)
 
 
-def standard_gamma(key, alpha, shape=(), dtype=np.float32):
+def standard_gamma(key, alpha, shape=(), dtype=np.float64):
+    dtype = xla_bridge.canonicalize_dtype(dtype)
     return _standard_gamma(key, alpha, shape, dtype)
 
 
 # TODO: inefficient implementation; jit currently fails due to
 # dynamic size of random.uniform.
 @partial(jit, static_argnums=(2, 3))
-def _binomial(key, p, n=1, shape=()):
+def _binomial(key, p, n, shape):
     p, n = promote_shapes(p, n)
     shape = shape or lax.broadcast_shapes(np.shape(p), np.shape(n))
     n_max = int(np.max(n))
@@ -216,7 +218,7 @@ def binomial(key, p, n=1, shape=()):
 
 
 @partial(jit, static_argnums=(2,))
-def _categorical(key, p, shape=()):
+def _categorical(key, p, shape):
     # this implementation is fast when event shape is small, and slow otherwise
     # Ref: https://stackoverflow.com/a/34190035
     shape = shape or p.shape[:-1]
@@ -232,7 +234,7 @@ def categorical(key, p, shape=()):
 
 
 @partial(jit, static_argnums=(2,))
-def _poisson(key, rate, shape=()):
+def _poisson(key, rate, shape):
     # Ref: https://en.wikipedia.org/wiki/Poisson_distribution#Generating_Poisson-distributed_random_variables
     shape = shape or np.shape(rate)
     L = np.exp(-rate)
diff --git a/numpyro/mcmc.py b/numpyro/mcmc.py
--- a/numpyro/mcmc.py
+++ b/numpyro/mcmc.py
@@ -48,8 +48,10 @@
 
 
 def _get_num_steps(step_size, trajectory_length):
-    num_steps = np.array(trajectory_length / step_size, dtype=np.int32)
-    return np.where(num_steps < 1, np.array(1, dtype=np.int32), num_steps)
+    num_steps = np.clip(trajectory_length / step_size, a_min=1)
+    # NB: casting to np.int64 does not take effect (returns np.int32 instead)
+    # if jax_enable_x64 is False
+    return num_steps.astype(np.int64)
 
 
 def _sample_momentum(unpack_fn, inverse_mass_matrix, rng):
diff --git a/setup.py b/setup.py
--- a/setup.py
+++ b/setup.py
@@ -11,7 +11,7 @@
     author='Uber AI Labs',
     author_email='npradhan@uber.com',
     install_requires=[
-        'jax>=0.1.26',
+        'jax>=0.1.35',
         'jaxlib>=0.1.14',
         'tqdm',
     ],


Test patch:
diff --git a/test/conftest.py b/test/conftest.py
--- a/test/conftest.py
+++ b/test/conftest.py
@@ -1,3 +1,5 @@
+import os
+
 from jax.config import config
 
 from numpyro.util import set_rng_seed
@@ -5,4 +7,6 @@
 
 def pytest_runtest_setup(item):
     config.update('jax_platform_name', 'cpu')
+    if 'JAX_ENABLE_x64' in os.environ:
+        config.update('jax_enable_x64', True)
     set_rng_seed(0)
diff --git a/test/test_mcmc.py b/test/test_mcmc.py
--- a/test/test_mcmc.py
+++ b/test/test_mcmc.py
@@ -1,3 +1,5 @@
+import os
+
 import numpy as onp
 import pytest
 from numpy.testing import assert_allclose
@@ -32,6 +34,9 @@ def potential_fn(z):
     assert_allclose(np.mean(hmc_states), true_mean, rtol=0.05)
     assert_allclose(np.std(hmc_states), true_std, rtol=0.05)
 
+    if 'JAX_ENABLE_x64' in os.environ:
+        assert hmc_states.dtype == np.float64
+
 
 @pytest.mark.parametrize('algo', ['HMC', 'NUTS'])
 def test_logistic_regression(algo):
@@ -48,9 +53,12 @@ def model(labels):
         return sample('obs', dist.Bernoulli(logits=logits), obs=labels)
 
     init_params, potential_fn, constrain_fn = initialize_model(random.PRNGKey(2), model, labels)
-    hmc_states = mcmc(warmup_steps, num_samples, init_params, sampler='hmc',
-                      potential_fn=potential_fn, trajectory_length=10, constrain_fn=constrain_fn)
-    assert_allclose(np.mean(hmc_states['coefs'], 0), true_coefs, atol=0.2)
+    samples = mcmc(warmup_steps, num_samples, init_params, sampler='hmc',
+                   potential_fn=potential_fn, trajectory_length=10, constrain_fn=constrain_fn)
+    assert_allclose(np.mean(samples['coefs'], 0), true_coefs, atol=0.21)
+
+    if 'JAX_ENABLE_x64' in os.environ:
+        assert samples['coefs'].dtype == np.float64
 
 
 @pytest.mark.parametrize('algo', ['HMC', 'NUTS'])
@@ -77,6 +85,9 @@ def model(data):
                               progbar=False)
     assert_allclose(np.mean(hmc_states['p_latent'], 0), true_probs, atol=0.05)
 
+    if 'JAX_ENABLE_x64' in os.environ:
+        assert hmc_states['p_latent'].dtype == np.float64
+
 
 @pytest.mark.parametrize('algo', ['HMC', 'NUTS'])
 def test_dirichlet_categorical(algo):
@@ -101,6 +112,9 @@ def model(data):
                               progbar=False)
     assert_allclose(np.mean(hmc_states['p_latent'], 0), true_probs, atol=0.02)
 
+    if 'JAX_ENABLE_x64' in os.environ:
+        assert hmc_states['p_latent'].dtype == np.float64
+
 
 def test_change_point():
     # Ref: https://forum.pyro.ai/t/i-dont-understand-why-nuts-code-is-not-working-bayesian-hackers-mail/696
@@ -122,8 +136,7 @@ def model(data):
         12,  35,  17,  23,  17,   4,   2,  31,  30,  13,  27,   0,  39,  37,
         5,  14,  13,  22,
     ])
-    init_params, potential_fn, constrain_fn = initialize_model(random.PRNGKey(2), model, count_data,
-                                                               init_strategy='prior')
+    init_params, potential_fn, constrain_fn = initialize_model(random.PRNGKey(4), model, count_data)
     init_kernel, sample_kernel = hmc(potential_fn)
     hmc_state = init_kernel(init_params, num_warmup=warmup_steps)
     hmc_states = fori_collect(num_samples, sample_kernel, hmc_state,
@@ -134,6 +147,11 @@ def model(data):
     mode = tau_values[mode_ind]
     assert mode == 44
 
+    if 'JAX_ENABLE_x64' in os.environ:
+        assert hmc_states['lambda1'].dtype == np.float64
+        assert hmc_states['lambda2'].dtype == np.float64
+        assert hmc_states['tau'].dtype == np.float64
+
 
 @pytest.mark.parametrize('with_logits', ['True', 'False'])
 def test_binomial_stable(with_logits):
@@ -156,3 +174,6 @@ def model(data):
                               transform=lambda x: constrain_fn(x.z))
 
     assert_allclose(np.mean(hmc_states['p'], 0), data['x'] / data['n'], rtol=0.05)
+
+    if 'JAX_ENABLE_x64' in os.environ:
+        assert hmc_states['p'].dtype == np.float64