Repository: pyro-ppl/numpyro

Issue:
fori_collect does not work with non-float arrays
The issue is ravel_pytree which calls `vjp` and does not work with int32 or bool arrays. As such, `fori_collect` will fail if the default identity transform is used.


Discussion:
@neerajprad  how about we patch [this function](https://github.com/google/jax/blob/master/jax/api.py#L952) to do nothing? it is used to raise errors if inputs of `vjp` are int/bool. Or we can create a `vjp` + `ravel_pytree` version which does not have this check. I prefer using `patch` but would like to hear your opinion too.

Resolving patch:
diff --git a/numpyro/util.py b/numpyro/util.py
--- a/numpyro/util.py
+++ b/numpyro/util.py
@@ -1,3 +1,4 @@
+from collections import namedtuple
 from contextlib import contextmanager
 import random
 
@@ -5,8 +6,8 @@
 import tqdm
 
 from jax import jit, lax, ops, vmap
-from jax.flatten_util import ravel_pytree
 import jax.numpy as np
+from jax.tree_util import tree_flatten, tree_unflatten, tree_map
 
 _DATA_TYPES = {}
 _DISABLE_CONTROL_FLOW_PRIM = False
@@ -171,3 +172,29 @@ def decorator(destin_class):
         return destin_class
 
     return decorator
+
+
+pytree_metadata = namedtuple('pytree_metadata', ['flat', 'shape', 'size', 'dtype'])
+
+
+def _ravel_list(*leaves):
+    leaves_metadata = tree_map(lambda l: pytree_metadata(np.ravel(l), np.shape(l), np.size(l), lax.dtype(l)),
+                               leaves)
+    leaves_idx = np.cumsum(np.array((0,) + tuple(d.size for d in leaves_metadata)))
+
+    def unravel_list(arr):
+        return [np.reshape(lax.dynamic_slice_in_dim(arr, leaves_idx[i], m.size),
+                           m.shape).astype(m.dtype)
+                for i, m in enumerate(leaves_metadata)]
+
+    return np.concatenate([m.flat for m in leaves_metadata]), unravel_list
+
+
+def ravel_pytree(pytree):
+    leaves, treedef = tree_flatten(pytree)
+    flat, unravel_list = _ravel_list(*leaves)
+
+    def unravel_pytree(arr):
+        return tree_unflatten(treedef, unravel_list(arr))
+
+    return flat, unravel_pytree


Test patch:
diff --git a/test/test_util.py b/test/test_util.py
--- a/test/test_util.py
+++ b/test/test_util.py
@@ -1,7 +1,12 @@
 import jax.numpy as np
+import pytest
+from jax import lax
+from jax.lib.xla_bridge import canonicalize_dtype
 from jax.test_util import check_eq
+from jax.tree_util import tree_multimap, tree_flatten
+from numpy.testing import assert_allclose
 
-from numpyro.util import fori_collect
+from numpyro.util import fori_collect, ravel_pytree
 
 
 def test_fori_collect():
@@ -12,3 +17,20 @@ def f(x):
     expected_tree = {'i': np.array([[0.], [2.]])}
     actual_tree = fori_collect(1, 3, f, a, transform=lambda a: {'i': a['i']})
     check_eq(actual_tree, expected_tree)
+
+
+@pytest.mark.parametrize('pytree', [
+    {'a': np.array(0.), 'b': np.array([[1., 2.], [3., 4.]])},
+    {'a': np.array(0), 'b': np.array([[1, 2], [3, 4]])},
+    {'a': np.array(0), 'b': np.array([[1., 2.], [3., 4.]])},
+    {'a': 0., 'b': np.array([[1., 2.], [3., 4.]])},
+    {'a': False, 'b': np.array([[1., 2.], [3., 4.]])},
+    [False, True, 0., np.array([[1., 2.], [3., 4.]])],
+])
+def test_ravel_pytree(pytree):
+    flat, unravel_fn = ravel_pytree(pytree)
+    unravel = unravel_fn(flat)
+    tree_flatten(tree_multimap(lambda x, y: assert_allclose(x, y), unravel, pytree))
+    assert all(tree_flatten(tree_multimap(lambda x, y:
+                                          canonicalize_dtype(lax.dtype(x)) == canonicalize_dtype(lax.dtype(y)),
+                                          unravel, pytree))[0])