<p dir="auto">Please:</p>
<ul class="contains-task-list">
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> Check for duplicate issues.</li>
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> Provide a complete example of how to reproduce the bug, wrapped in triple backticks like this:</li>
</ul>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="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 &lt; 5.3
print(y.dtype)          # =&gt; &lt;dtype: 'bool'&gt;
print(unravel(y).dtype) # =&gt; &lt;dtype: 'float32'&gt;"><pre class="notranslate"><span class="pl-k">import</span> <span class="pl-s1">jax</span>.<span class="pl-s1">numpy</span> <span class="pl-k">as</span> <span class="pl-s1">jnp</span>
<span class="pl-k">from</span> <span class="pl-s1">jax</span>.<span class="pl-s1">flatten_util</span> <span class="pl-k">import</span> <span class="pl-s1">ravel_pytree</span>
<span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-s1">jnp</span>.<span class="pl-en">arange</span>(<span class="pl-c1">10</span>, <span class="pl-s1">dtype</span><span class="pl-c1">=</span><span class="pl-s1">jnp</span>.<span class="pl-s1">float32</span>)
<span class="pl-s1">x_flat</span>, <span class="pl-s1">unravel</span> <span class="pl-c1">=</span> <span class="pl-en">ravel_pytree</span>(<span class="pl-s1">x</span>)
<span class="pl-s1">y</span> <span class="pl-c1">=</span> <span class="pl-s1">x_flat</span> <span class="pl-c1">&lt;</span> <span class="pl-c1">5.3</span>
<span class="pl-en">print</span>(<span class="pl-s1">y</span>.<span class="pl-s1">dtype</span>)          <span class="pl-c"># =&gt; &lt;dtype: 'bool'&gt;</span>
<span class="pl-en">print</span>(<span class="pl-en">unravel</span>(<span class="pl-s1">y</span>).<span class="pl-s1">dtype</span>) <span class="pl-c"># =&gt; &lt;dtype: 'float32'&gt;</span></pre></div>
<ul class="contains-task-list">
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> If applicable, include full error messages/tracebacks.<br>
n/a</li>
</ul>
<p dir="auto">I'm running Python 3.9 with jax 0.2.19 and jaxlib 0.1.70.</p>
<p dir="auto">Please:</p>
<ul class="contains-task-list">
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> Check for duplicate issues.</li>
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> Provide a complete example of how to reproduce the bug, wrapped in triple backticks like this:</li>
</ul>
<p dir="auto">Steps to reproduce:</p>
<ul dir="auto">
<li>Clone <a class="commit-link" data-hovercard-type="commit" data-hovercard-url="https://github.com/NixOS/nixpkgs/commit/19c42cefd72befdbeab5b98c5534626d5905a8ce/hovercard" href="https://github.com/NixOS/nixpkgs/commit/19c42cefd72befdbeab5b98c5534626d5905a8ce">NixOS/nixpkgs@<tt>19c42ce</tt></a></li>
<li>Update <code class="notranslate">pkgs/development/python-modules/jaxlib/default.nix</code> with these contents <a href="https://gist.github.com/samuela/931d7e29dd5e83ec5bf73de3bbc5b084">https://gist.github.com/samuela/931d7e29dd5e83ec5bf73de3bbc5b084</a>.</li>
<li>Run <code class="notranslate">NIX_PATH=.. nix-build -E "with import &lt;nixpkgs&gt; {}; python3Packages.jaxlib.override { cudaSupport = true; }"</code> inside your <code class="notranslate">nixpkgs</code> checkout.</li>
</ul>
<p dir="auto">Observe the failure:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="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&lt;class ... Args&gt; using FormatSpec = absl::lts_20211102::str_format_internal::FormatSpecTemplate&lt;(ArgumentToConv&lt;Args&gt;)()...&gt; [with Args = {google::protobuf::stringpiece_internal::StringPiece}]’
external/org_tensorflow/tensorflow/compiler/xla/util.h:214:8:   required by substitution of ‘template&lt;class ... Args&gt; tensorflow::Status xla::InvalidArgument(absl::lts_20211102::FormatSpec&lt;Args ...&gt;&amp;, const Args&amp; ...) [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&amp;, const absl::lts_20211102::str_format_internal::FormatConversionSpecImpl&amp;, absl::lts_20211102::str_format_internal::FormatSinkImpl*)’
  305 |       decltype(str_format_internal::FormatConvertImpl(
      |                ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^
  306 |           std::declval&lt;const Arg&amp;&gt;(),
      |           ~~~~~~~~~~~~~~~~~~~~~~~~~~~
  307 |           std::declval&lt;const FormatConversionSpecImpl&amp;&gt;(),
      |           ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
  308 |           std::declval&lt;FormatSinkImpl*&gt;())){});
      |           ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
external/com_google_absl/absl/strings/internal/str_format/arg.h:60:6: note: candidate: ‘template&lt;class T&gt; decltype (absl::lts_20211102::str_format_internal::AbslFormatConvert(v, declval&lt;const absl::lts_20211102::FormatConversionSpec&amp;&gt;(), declval&lt;absl::lts_20211102::FormatSink*&gt;())) absl::lts_20211102::str_format_internal::FormatConvertImpl(const T&amp;, absl::lts_20211102::str_format_internal::FormatConversionSpecImpl, absl::lts_20211102::str_format_internal::FormatSinkImpl*)’
   60 | auto FormatConvertImpl(const T&amp; v, FormatConversionSpecImpl conv,
      |      ^~~~~~~~~~~~~~~~~"><pre class="notranslate"><code class="notranslate">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&lt;class ... Args&gt; using FormatSpec = absl::lts_20211102::str_format_internal::FormatSpecTemplate&lt;(ArgumentToConv&lt;Args&gt;)()...&gt; [with Args = {google::protobuf::stringpiece_internal::StringPiece}]’
external/org_tensorflow/tensorflow/compiler/xla/util.h:214:8:   required by substitution of ‘template&lt;class ... Args&gt; tensorflow::Status xla::InvalidArgument(absl::lts_20211102::FormatSpec&lt;Args ...&gt;&amp;, const Args&amp; ...) [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&amp;, const absl::lts_20211102::str_format_internal::FormatConversionSpecImpl&amp;, absl::lts_20211102::str_format_internal::FormatSinkImpl*)’
  305 |       decltype(str_format_internal::FormatConvertImpl(
      |                ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^
  306 |           std::declval&lt;const Arg&amp;&gt;(),
      |           ~~~~~~~~~~~~~~~~~~~~~~~~~~~
  307 |           std::declval&lt;const FormatConversionSpecImpl&amp;&gt;(),
      |           ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
  308 |           std::declval&lt;FormatSinkImpl*&gt;())){});
      |           ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
external/com_google_absl/absl/strings/internal/str_format/arg.h:60:6: note: candidate: ‘template&lt;class T&gt; decltype (absl::lts_20211102::str_format_internal::AbslFormatConvert(v, declval&lt;const absl::lts_20211102::FormatConversionSpec&amp;&gt;(), declval&lt;absl::lts_20211102::FormatSink*&gt;())) absl::lts_20211102::str_format_internal::FormatConvertImpl(const T&amp;, absl::lts_20211102::str_format_internal::FormatConversionSpecImpl, absl::lts_20211102::str_format_internal::FormatSinkImpl*)’
   60 | auto FormatConvertImpl(const T&amp; v, FormatConversionSpecImpl conv,
      |      ^~~~~~~~~~~~~~~~~
</code></pre></div>
<p dir="auto">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 <code class="notranslate">nix-build -A python3Packages.jaxlib</code>) runs just fine.</p>
<ul class="contains-task-list">
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> If applicable, include full error messages/tracebacks.</li>
</ul>
<p dir="auto">You can check out the unabridged build log here: <a href="https://gist.github.com/samuela/fca69ea5ad0eb20d77b8be060446db21">https://gist.github.com/samuela/fca69ea5ad0eb20d77b8be060446db21</a>.</p>