Repository: google/jax

Issue:
Unraveling an array with dtype bool results in pytree arrays with dtype float32
Please:

- [x] Check for duplicate issues.
- [x] Provide a complete example of how to reproduce the bug, wrapped in triple backticks like this:

```python
import jax.numpy as jnp
from jax.flatten_util import ravel_pytree
x = jnp.arange(10, dtype=jnp.float32)
x_flat, unravel = ravel_pytree(x)
y = x_flat < 5.3
print(y.dtype)          # => <dtype: 'bool'>
print(unravel(y).dtype) # => <dtype: 'float32'>
```

- [x] If applicable, include full error messages/tracebacks.
n/a

I'm running Python 3.9 with jax 0.2.19 and jaxlib 0.1.70.


Discussion:
Hi - thanks for the report. I think this may be intended, given the lines here: https://github.com/google/jax/blob/c69eedd6f8bd360f23b8793c074648addd1e23d7/jax/flatten_util.py#L55-L56

The output dtype is explicitly set to the type promotion output of the input dtypes.

Maybe @mattjj has some context on this?
I was worried about that... My hope was that ravel/unravel would only transform array shapes and not dtypes, but perhaps there a good reason it should coerce dtypes as well?
Sorry @samuela, I let this one slip through the cracks somehow!

The generated `unravel` function just isn't dtype-polymorphic; it produces an output with type (i.e. pytree structure, array shapes, and array dtypes) equal to that of the provided input. Perhaps it should've even raised an error when given a bool-dtype input...

One reason for that is to handle the case when there are mixed dtypes in the input pytree. Since the raveled array needs to have a uniform dtype, we promote the dtypes of the inputs, and when unraveling we cast back to the input's dtypes. But that means the output dtype doesn't determine the input dtypes in general, and hence it's not clear how to make an `unravel` which is polymorphic in the way you expected.

I'm not sure how to reconcile these two behaviors (supporting multiple distinct input dtypes, and producing a dtype-polymorphic `unravel`) in one function, without making the behavior complicated. We could just special case it to be "if all the input dtypes are equal, then produce a polymorphic `unravel`". But that might be making a simple function kind of unpredictable.

Since the function is small, would it make sense just to have your own implementation? (Or, in the intervening time, did you find other solutions?)

Here's an implementation which has the behavior you expected:

```python
import jax
from jax.tree_util import tree_flatten, tree_unflatten
from jax.util import unzip2
import jax.numpy as jnp
import numpy as np

def ravel_pytree(pytree):
  leaves, treedef = tree_flatten(pytree)
  flat, unravel_list = _ravel_list(leaves)
  unravel_pytree = lambda flat: tree_unflatten(treedef, unravel_list(flat))
  return flat, unravel_pytree

def _ravel_list(lst):
  if not lst: return jnp.array([], jnp.float32), lambda _: []

  from_dtypes = [jax.dtypes.result_type(l) for l in lst]
  to_dtype = jax.dtypes.result_type(*from_dtypes)
  if not all(from_dtype == to_dtype for from_dtype in from_dtypes):
    raise Exception
  del from_dtypes, to_dtype

  sizes, shapes = unzip2((jnp.size(x), jnp.shape(x)) for x in lst)
  indices = np.cumsum(sizes)

  def unravel(arr):
    chunks = jnp.split(arr, indices[:-1])
    return [chunk.reshape(shape) for chunk, shape in zip(chunks, shapes)]

  raveled = jnp.concatenate([jnp.ravel(e) for e in lst])
  return raveled, unravel


x = jnp.arange(10, dtype=jnp.float32)
x_flat, unravel = ravel_pytree(x)
y = x_flat < 5.3
print(y.dtype)          # => <dtype: 'bool'>
print(unravel(y).dtype) # => <dtype: 'bool'>
```

Since I'm in an issue-closing mood, and since this is an old one, I'm going to somewhat preemptively close this issue. Let us know if we should reopen and continue the discussion!

Resolving patch:
diff --git a/jax/_src/flatten_util.py b/jax/_src/flatten_util.py
--- a/jax/_src/flatten_util.py
+++ b/jax/_src/flatten_util.py
@@ -57,7 +57,23 @@ def _ravel_list(lst):
   sizes, shapes = unzip2((jnp.size(x), jnp.shape(x)) for x in lst)
   indices = np.cumsum(sizes)
 
+  if all(dt == to_dtype for dt in from_dtypes):
+    # Skip any dtype conversion, resulting in a dtype-polymorphic `unravel`.
+    # See https://github.com/google/jax/issues/7809.
+    del from_dtypes, to_dtype
+    def unravel(arr):
+      chunks = jnp.split(arr, indices[:-1])
+      return [chunk.reshape(shape) for chunk, shape in zip(chunks, shapes)]
+    raveled = jnp.concatenate([jnp.ravel(e) for e in lst])
+    return raveled, unravel
+
+  # When there is more than one distinct input dtype, we perform type
+  # conversions and produce a dtype-specific unravel function.
   def unravel(arr):
+    arr_dtype = dtypes.dtype(arr)
+    if arr_dtype != to_dtype:
+      raise TypeError(f"unravel function given array of dtype {arr_dtype}, "
+                      f"but expected dtype {to_dtype}")
     chunks = jnp.split(arr, indices[:-1])
     with warnings.catch_warnings():
       warnings.simplefilter("ignore")  # ignore complex-to-real cast warning


Test patch:
diff --git a/tests/tree_util_test.py b/tests/tree_util_test.py
--- a/tests/tree_util_test.py
+++ b/tests/tree_util_test.py
@@ -379,6 +379,23 @@ def testEmpty(self):
     tree_ = unravel(raveled)
     self.assertAllClose(tree, tree_, atol=0., rtol=0.)
 
+  def testDtypePolymorphicUnravel(self):
+    # https://github.com/google/jax/issues/7809
+    x = jnp.arange(10, dtype=jnp.float32)
+    x_flat, unravel = flatten_util.ravel_pytree(x)
+    y = x_flat < 5.3
+    x_ = unravel(y)
+    self.assertEqual(x_.dtype, y.dtype)
+
+  def testDtypeMonomorphicUnravel(self):
+    # https://github.com/google/jax/issues/7809
+    x1 = jnp.arange(10, dtype=jnp.float32)
+    x2 = jnp.arange(10, dtype=jnp.int32)
+    x_flat, unravel = flatten_util.ravel_pytree((x1, x2))
+    y = x_flat < 5.3
+    with self.assertRaisesRegex(TypeError, 'but expected dtype'):
+      _ = unravel(y)
+
 
 if __name__ == "__main__":
   absltest.main(testLoader=jtu.JaxTestLoader())