Issue 1:
Please:

Check for duplicate issues.

Provide a complete example of how to reproduce the bug, wrapped in triple backticks like this:

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'>

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.

Issue 2:
Please:

Check for duplicate issues.

Provide a complete example of how to reproduce the bug, wrapped in triple backticks like this:

Steps to reproduce:

Clone NixOS/nixpkgs@19c42ce

Update pkgs/development/python-modules/jaxlib/default.nix with these contents https://gist.github.com/samuela/931d7e29dd5e83ec5bf73de3bbc5b084.

Run NIX_PATH=.. nix-build -E "with import <nixpkgs> {}; python3Packages.jaxlib.override { cudaSupport = true; }" inside your nixpkgs checkout.

Observe the failure:

ERROR: /build/output/external/org_tensorflow/tensorflow/compiler/xla/python/BUILD:206:11: Compiling tensorflow/compiler/xla/python/pprof_profile_builder.cc failed: (Exit 1): crosstool_wrapper_driver_is_not_gcc failed: error executing command external/local_config_cuda/crosstool/clang/bin/crosstool_wrapper_driver_is_not_gcc -MD -MF ... (remaining 184 arguments skipped)
In file included from external/com_google_absl/absl/strings/str_format.h:76,
from external/org_tensorflow/tensorflow/compiler/xla/util.h:35,
from external/org_tensorflow/tensorflow/compiler/xla/python/pprof_profile_builder.cc:21:
external/com_google_absl/absl/strings/internal/str_format/arg.h: In instantiation of ‘constexpr absl::lts_20211102::FormatConversionCharSet absl::lts_20211102::str_format_internal::ArgumentToConv() [with Arg = google::protobuf::stringpiece_internal::StringPiece]’:
external/com_google_absl/absl/strings/str_format.h:257:46: required by substitution of ‘template<class ... Args> using FormatSpec = absl::lts_20211102::str_format_internal::FormatSpecTemplate<(ArgumentToConv<Args>)()...> [with Args = {google::protobuf::stringpiece_internal::StringPiece}]’
external/org_tensorflow/tensorflow/compiler/xla/util.h:214:8: required by substitution of ‘template<class ... Args> tensorflow::Status xla::InvalidArgument(absl::lts_20211102::FormatSpec<Args ...>&, const Args& ...) [with Args = {google::protobuf::stringpiece_internal::StringPiece}]’
external/org_tensorflow/tensorflow/compiler/xla/python/pprof_profile_builder.cc:71:71: required from here
external/com_google_absl/absl/strings/internal/str_format/arg.h:305:54: error: no matching function for call to ‘FormatConvertImpl(const google::protobuf::stringpiece_internal::StringPiece&, const absl::lts_20211102::str_format_internal::FormatConversionSpecImpl&, absl::lts_20211102::str_format_internal::FormatSinkImpl*)’
305 | decltype(str_format_internal::FormatConvertImpl(
| ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^
306 | std::declval<const Arg&>(),
| ~~~~~~~~~~~~~~~~~~~~~~~~~~~
307 | std::declval<const FormatConversionSpecImpl&>(),
| ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
308 | std::declval<FormatSinkImpl*>())){});
| ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
external/com_google_absl/absl/strings/internal/str_format/arg.h:60:6: note: candidate: ‘template<class T> decltype (absl::lts_20211102::str_format_internal::AbslFormatConvert(v, declval<const absl::lts_20211102::FormatConversionSpec&>(), declval<absl::lts_20211102::FormatSink*>())) absl::lts_20211102::str_format_internal::FormatConvertImpl(const T&, absl::lts_20211102::str_format_internal::FormatConversionSpecImpl, absl::lts_20211102::str_format_internal::FormatSinkImpl*)’
60 | auto FormatConvertImpl(const T& v, FormatConversionSpecImpl conv,
| ^~~~~~~~~~~~~~~~~

I can confirm that this build "script" works on 0.1.75 and a number versions before that. Oddly enough the CPU-only build (run nix-build -A python3Packages.jaxlib) runs just fine.

If applicable, include full error messages/tracebacks.

You can check out the unabridged build log here: https://gist.github.com/samuela/fca69ea5ad0eb20d77b8be060446db21.