<p dir="auto">Please:</p>
<ul dir="auto">
<li>[ x] Check for duplicate issues.</li>
<li>[ x] Provide a complete example of how to reproduce the bug, wrapped in triple backticks like this:<br>
I am installing the jax using the following command:<br>
'pip install --upgrade "jax[cuda111]" jaxlib==0.1.67 -f <a href="https://storage.googleapis.com/jax-releases/jax_releases.html" rel="nofollow">https://storage.googleapis.com/jax-releases/jax_releases.html</a>'<br>
it gives me the following error:<br>
jax depends on jaxlib@ <a href="https://storage.googleapis.com/jax-releases/cuda111/jaxlib-0.1.67+cuda111-cp38-none-manylinux2010_x86_64.whl" rel="nofollow">https://storage.googleapis.com/jax-releases/cuda111/jaxlib-0.1.67+cuda111-cp38-none-manylinux2010_x86_64.whl</a><br>
I am using cuda 11.4 with python 3.8.10</li>
</ul>
<p dir="auto">it give</p>
<ul class="contains-task-list">
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox"> If applicable, include full error messages/tracebacks.</li>
</ul>
<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 itertools

import jax
import jax.numpy as jnp

import haiku as hk
import optax


Model = hk.nets.ResNet152


def dataloader(batch_size, *, key):
    while True:
        # Binary classification
        X_key, Y_key = jax.random.split(key)
        X = jax.random.normal(X_key, (batch_size, 64, 64, 3))
        Y = jax.random.randint(Y_key, (batch_size,), 0, 2)
        yield X, Y


def train():
    batch_size = 32
    learning_rate = 0.001

    def forward(input, is_training):
        net = Model(num_classes=2)
        output = net(input, is_training)
        return output

    def learner_fn(input, ground_truth):
        logits = forward(input, True)
        labels = jax.nn.one_hot(ground_truth, 2)
        return jnp.mean(optax.softmax_cross_entropy(logits, labels))

    learner_fn_t = hk.transform_with_state(learner_fn)
    learner_fn_t = hk.without_apply_rng(learner_fn_t)

    @jax.jit
    def train_step(params, state, opt_state, X, y):
        (loss, state), grads = jax.value_and_grad(learner_fn_t.apply, has_aux=True)(params, state, X, y)
        updates, opt_state = optimizer.update(grads, opt_state, params)
        params = optax.apply_updates(params, updates)
        return params, state, opt_state, loss
    
    rng = jax.random.PRNGKey(42)
    loader_rng, rng = jax.random.split(rng)
    loader = dataloader(batch_size, key=loader_rng)

    X, y = next(loader)
    params, state = learner_fn_t.init(rng, X, y)

    optimizer = optax.adam(learning_rate)
    opt_state = optimizer.init(params)

    # Training
    for X, y in itertools.islice(loader, 50):
        params, state, opt_state, loss = train_step(params, state, opt_state, X, y)
        print(loss)


if __name__ == '__main__':
    train()"><pre class="notranslate"><span class="pl-k">import</span> <span class="pl-s1">itertools</span>

<span class="pl-k">import</span> <span class="pl-s1">jax</span>
<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">import</span> <span class="pl-s1">haiku</span> <span class="pl-k">as</span> <span class="pl-s1">hk</span>
<span class="pl-k">import</span> <span class="pl-s1">optax</span>


<span class="pl-v">Model</span> <span class="pl-c1">=</span> <span class="pl-s1">hk</span>.<span class="pl-s1">nets</span>.<span class="pl-v">ResNet152</span>


<span class="pl-k">def</span> <span class="pl-en">dataloader</span>(<span class="pl-s1">batch_size</span>, <span class="pl-c1">*</span>, <span class="pl-s1">key</span>):
    <span class="pl-k">while</span> <span class="pl-c1">True</span>:
        <span class="pl-c"># Binary classification</span>
        <span class="pl-v">X_key</span>, <span class="pl-v">Y_key</span> <span class="pl-c1">=</span> <span class="pl-s1">jax</span>.<span class="pl-s1">random</span>.<span class="pl-en">split</span>(<span class="pl-s1">key</span>)
        <span class="pl-v">X</span> <span class="pl-c1">=</span> <span class="pl-s1">jax</span>.<span class="pl-s1">random</span>.<span class="pl-en">normal</span>(<span class="pl-v">X_key</span>, (<span class="pl-s1">batch_size</span>, <span class="pl-c1">64</span>, <span class="pl-c1">64</span>, <span class="pl-c1">3</span>))
        <span class="pl-v">Y</span> <span class="pl-c1">=</span> <span class="pl-s1">jax</span>.<span class="pl-s1">random</span>.<span class="pl-en">randint</span>(<span class="pl-v">Y_key</span>, (<span class="pl-s1">batch_size</span>,), <span class="pl-c1">0</span>, <span class="pl-c1">2</span>)
        <span class="pl-k">yield</span> <span class="pl-v">X</span>, <span class="pl-v">Y</span>


<span class="pl-k">def</span> <span class="pl-en">train</span>():
    <span class="pl-s1">batch_size</span> <span class="pl-c1">=</span> <span class="pl-c1">32</span>
    <span class="pl-s1">learning_rate</span> <span class="pl-c1">=</span> <span class="pl-c1">0.001</span>

    <span class="pl-k">def</span> <span class="pl-en">forward</span>(<span class="pl-s1">input</span>, <span class="pl-s1">is_training</span>):
        <span class="pl-s1">net</span> <span class="pl-c1">=</span> <span class="pl-v">Model</span>(<span class="pl-s1">num_classes</span><span class="pl-c1">=</span><span class="pl-c1">2</span>)
        <span class="pl-s1">output</span> <span class="pl-c1">=</span> <span class="pl-en">net</span>(<span class="pl-s1">input</span>, <span class="pl-s1">is_training</span>)
        <span class="pl-k">return</span> <span class="pl-s1">output</span>

    <span class="pl-k">def</span> <span class="pl-en">learner_fn</span>(<span class="pl-s1">input</span>, <span class="pl-s1">ground_truth</span>):
        <span class="pl-s1">logits</span> <span class="pl-c1">=</span> <span class="pl-en">forward</span>(<span class="pl-s1">input</span>, <span class="pl-c1">True</span>)
        <span class="pl-s1">labels</span> <span class="pl-c1">=</span> <span class="pl-s1">jax</span>.<span class="pl-s1">nn</span>.<span class="pl-en">one_hot</span>(<span class="pl-s1">ground_truth</span>, <span class="pl-c1">2</span>)
        <span class="pl-k">return</span> <span class="pl-s1">jnp</span>.<span class="pl-en">mean</span>(<span class="pl-s1">optax</span>.<span class="pl-en">softmax_cross_entropy</span>(<span class="pl-s1">logits</span>, <span class="pl-s1">labels</span>))

    <span class="pl-s1">learner_fn_t</span> <span class="pl-c1">=</span> <span class="pl-s1">hk</span>.<span class="pl-en">transform_with_state</span>(<span class="pl-s1">learner_fn</span>)
    <span class="pl-s1">learner_fn_t</span> <span class="pl-c1">=</span> <span class="pl-s1">hk</span>.<span class="pl-en">without_apply_rng</span>(<span class="pl-s1">learner_fn_t</span>)

    <span class="pl-en">@<span class="pl-s1">jax</span>.<span class="pl-s1">jit</span></span>
    <span class="pl-k">def</span> <span class="pl-en">train_step</span>(<span class="pl-s1">params</span>, <span class="pl-s1">state</span>, <span class="pl-s1">opt_state</span>, <span class="pl-v">X</span>, <span class="pl-s1">y</span>):
        (<span class="pl-s1">loss</span>, <span class="pl-s1">state</span>), <span class="pl-s1">grads</span> <span class="pl-c1">=</span> <span class="pl-s1">jax</span>.<span class="pl-en">value_and_grad</span>(<span class="pl-s1">learner_fn_t</span>.<span class="pl-s1">apply</span>, <span class="pl-s1">has_aux</span><span class="pl-c1">=</span><span class="pl-c1">True</span>)(<span class="pl-s1">params</span>, <span class="pl-s1">state</span>, <span class="pl-v">X</span>, <span class="pl-s1">y</span>)
        <span class="pl-s1">updates</span>, <span class="pl-s1">opt_state</span> <span class="pl-c1">=</span> <span class="pl-s1">optimizer</span>.<span class="pl-en">update</span>(<span class="pl-s1">grads</span>, <span class="pl-s1">opt_state</span>, <span class="pl-s1">params</span>)
        <span class="pl-s1">params</span> <span class="pl-c1">=</span> <span class="pl-s1">optax</span>.<span class="pl-en">apply_updates</span>(<span class="pl-s1">params</span>, <span class="pl-s1">updates</span>)
        <span class="pl-k">return</span> <span class="pl-s1">params</span>, <span class="pl-s1">state</span>, <span class="pl-s1">opt_state</span>, <span class="pl-s1">loss</span>
    
    <span class="pl-s1">rng</span> <span class="pl-c1">=</span> <span class="pl-s1">jax</span>.<span class="pl-s1">random</span>.<span class="pl-v">PRNGKey</span>(<span class="pl-c1">42</span>)
    <span class="pl-s1">loader_rng</span>, <span class="pl-s1">rng</span> <span class="pl-c1">=</span> <span class="pl-s1">jax</span>.<span class="pl-s1">random</span>.<span class="pl-en">split</span>(<span class="pl-s1">rng</span>)
    <span class="pl-s1">loader</span> <span class="pl-c1">=</span> <span class="pl-en">dataloader</span>(<span class="pl-s1">batch_size</span>, <span class="pl-s1">key</span><span class="pl-c1">=</span><span class="pl-s1">loader_rng</span>)

    <span class="pl-v">X</span>, <span class="pl-s1">y</span> <span class="pl-c1">=</span> <span class="pl-en">next</span>(<span class="pl-s1">loader</span>)
    <span class="pl-s1">params</span>, <span class="pl-s1">state</span> <span class="pl-c1">=</span> <span class="pl-s1">learner_fn_t</span>.<span class="pl-en">init</span>(<span class="pl-s1">rng</span>, <span class="pl-v">X</span>, <span class="pl-s1">y</span>)

    <span class="pl-s1">optimizer</span> <span class="pl-c1">=</span> <span class="pl-s1">optax</span>.<span class="pl-en">adam</span>(<span class="pl-s1">learning_rate</span>)
    <span class="pl-s1">opt_state</span> <span class="pl-c1">=</span> <span class="pl-s1">optimizer</span>.<span class="pl-en">init</span>(<span class="pl-s1">params</span>)

    <span class="pl-c"># Training</span>
    <span class="pl-k">for</span> <span class="pl-v">X</span>, <span class="pl-s1">y</span> <span class="pl-c1">in</span> <span class="pl-s1">itertools</span>.<span class="pl-en">islice</span>(<span class="pl-s1">loader</span>, <span class="pl-c1">50</span>):
        <span class="pl-s1">params</span>, <span class="pl-s1">state</span>, <span class="pl-s1">opt_state</span>, <span class="pl-s1">loss</span> <span class="pl-c1">=</span> <span class="pl-en">train_step</span>(<span class="pl-s1">params</span>, <span class="pl-s1">state</span>, <span class="pl-s1">opt_state</span>, <span class="pl-v">X</span>, <span class="pl-s1">y</span>)
        <span class="pl-en">print</span>(<span class="pl-s1">loss</span>)


<span class="pl-k">if</span> <span class="pl-s1">__name__</span> <span class="pl-c1">==</span> <span class="pl-s">'__main__'</span>:
    <span class="pl-en">train</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.</li>
</ul>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="$ XLA_FLAGS=--xla_dump_to=/tmp/slow_operation_dump pipenv run python3 slow_operation.py
WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
2022-05-21 01:56:39.688109: E external/org_tensorflow/tensorflow/compiler/xla/service/slow_operation_alarm.cc:61]
********************************
Very slow compile?  If you want to file a bug, run with envvar XLA_FLAGS=--xla_dump_to=/tmp/foo and attach the results.
Compiling module jit_train_step.261
********************************"><pre class="notranslate"><code class="notranslate">$ XLA_FLAGS=--xla_dump_to=/tmp/slow_operation_dump pipenv run python3 slow_operation.py
WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
2022-05-21 01:56:39.688109: E external/org_tensorflow/tensorflow/compiler/xla/service/slow_operation_alarm.cc:61]
********************************
Very slow compile?  If you want to file a bug, run with envvar XLA_FLAGS=--xla_dump_to=/tmp/foo and attach the results.
Compiling module jit_train_step.261
********************************
</code></pre></div>
<p dir="auto">Here is the content of <code class="notranslate">/tmp/slow_operation_dump</code>: <a href="https://github.com/google/jax/files/8747119/slow_operation_dump.zip">slow_operation_dump.zip</a>. Additionally, here is the Pipenv lock files to help you reproduce my development environment on macOS 12.3.1: <a href="https://github.com/google/jax/files/8747124/Pipfile.zip">Pipfile.zip</a></p>