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.
diff --git a/jax/_src/flatten_util.py b/jax/_src/flatten_util.py
index 2f7bce645..b7d7dc5b8 100644
--- 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
diff --git a/tests/tree_util_test.py b/tests/tree_util_test.py
index 0164f6d8a..0fd2ded31 100644
--- a/tests/tree_util_test.py
+++ b/tests/tree_util_test.py
@@ -379,6 +379,23 @@ class RavelUtilTest(jtu.JaxTestCase):
     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())
['tests/tree_util_test.py', 'jax/_src/flatten_util.py']