<p dir="auto">I wanted to have two <code class="notranslate">pmap()</code>-ed function, piping one's output into another. For example, in the following way:</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="for i, epoch in enumerate(tqdm(range(1, num_train_epochs + 1), desc=f&quot;Epoch ...&quot;, position=0, leave=True)):
    rng, input_rng = jax.random.split(rng)

    # train
    for batch in glue_train_data_loader(input_rng, train_dataset, total_batch_size):
        state, train_metrics, dropout_rngs = parallel_train_step(state, batch, dropout_rngs)

    # evaluate
    for batch in glue_eval_data_loader(eval_dataset, total_batch_size):
          predictions = parallel_eval_step(state, batch)
          metric.add_batch(predictions=chain(*predictions), references=chain(*labels))

    eval_metric = metric.compute()"><pre class="notranslate"><span class="pl-k">for</span> <span class="pl-s1">i</span>, <span class="pl-s1">epoch</span> <span class="pl-c1">in</span> <span class="pl-en">enumerate</span>(<span class="pl-en">tqdm</span>(<span class="pl-en">range</span>(<span class="pl-c1">1</span>, <span class="pl-s1">num_train_epochs</span> <span class="pl-c1">+</span> <span class="pl-c1">1</span>), <span class="pl-s1">desc</span><span class="pl-c1">=</span><span class="pl-s">f"Epoch ..."</span>, <span class="pl-s1">position</span><span class="pl-c1">=</span><span class="pl-c1">0</span>, <span class="pl-s1">leave</span><span class="pl-c1">=</span><span class="pl-c1">True</span>)):
    <span class="pl-s1">rng</span>, <span class="pl-s1">input_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-c"># train</span>
    <span class="pl-k">for</span> <span class="pl-s1">batch</span> <span class="pl-c1">in</span> <span class="pl-en">glue_train_data_loader</span>(<span class="pl-s1">input_rng</span>, <span class="pl-s1">train_dataset</span>, <span class="pl-s1">total_batch_size</span>):
        <span class="pl-s1">state</span>, <span class="pl-s1">train_metrics</span>, <span class="pl-s1">dropout_rngs</span> <span class="pl-c1">=</span> <span class="pl-en">parallel_train_step</span>(<span class="pl-s1">state</span>, <span class="pl-s1">batch</span>, <span class="pl-s1">dropout_rngs</span>)

    <span class="pl-c"># evaluate</span>
    <span class="pl-k">for</span> <span class="pl-s1">batch</span> <span class="pl-c1">in</span> <span class="pl-en">glue_eval_data_loader</span>(<span class="pl-s1">eval_dataset</span>, <span class="pl-s1">total_batch_size</span>):
          <span class="pl-s1">predictions</span> <span class="pl-c1">=</span> <span class="pl-en">parallel_eval_step</span>(<span class="pl-s1">state</span>, <span class="pl-s1">batch</span>)
          <span class="pl-s1">metric</span>.<span class="pl-en">add_batch</span>(<span class="pl-s1">predictions</span><span class="pl-c1">=</span><span class="pl-en">chain</span>(<span class="pl-c1">*</span><span class="pl-s1">predictions</span>), <span class="pl-s1">references</span><span class="pl-c1">=</span><span class="pl-en">chain</span>(<span class="pl-c1">*</span><span class="pl-s1">labels</span>))

    <span class="pl-s1">eval_metric</span> <span class="pl-c1">=</span> <span class="pl-s1">metric</span>.<span class="pl-en">compute</span>()</pre></div>
<p dir="auto"><code class="notranslate">parallel_train_step()</code> and <code class="notranslate">parallel_eval_step()</code> are the two <code class="notranslate">pmap()</code>-ed functions and <code class="notranslate">state</code> is the object being passed from one into another. The above code is adapted from <a href="https://colab.research.google.com/github/huggingface/notebooks/blob/master/examples/text_classification_flax.ipynb#scrollTo=FBARxJ8soOrs" rel="nofollow">this</a> HuggingFace Notebook.</p>
<p dir="auto">Now, in case of small models, the <code class="notranslate">state</code> is small in size and the output of <code class="notranslate">parallel_train_step</code> or <code class="notranslate">parallel_eval_step</code> gets stored in a buffer distinct from the input buffer. But in case of large models, the output <code class="notranslate">state</code> may get stored in the input buffer itself. But the input buffer is marked as donated, so when the output <code class="notranslate">state</code> is passed into the second function(<code class="notranslate">parallel_eval_step()</code>), it raises <code class="notranslate">RuntimeError: Invalid argument: Buffer has been deleted or donated</code> .</p>
<p dir="auto">I notice that we can <code class="notranslate">print()</code> the returned <code class="notranslate">state</code>, we can also <code class="notranslate">serialize()</code> it, but can't pass the buffers to the second <code class="notranslate">pmap()</code>-ed function.</p>
<p dir="auto">Manually, I tried solving the above error by using <code class="notranslate">state = jax.tree_util.tree_map(lambda x: jnp.array(x.tolist()), state)</code> before passing it to the second function. Also, tried using <code class="notranslate">lambda x: jnp.array(x,copy=True)</code>. Both of which didn't work. Is there some other way to solve the problem?</p>
<p dir="auto">Possible solution can be to un-mark the buffer as <code class="notranslate">donated</code> when the input buffer is used to store the output. But how to do it?</p>
<p dir="auto"><strong>Related:</strong> <a class="issue-link js-issue-link" data-error-text="Failed to load title" data-id="526635040" data-permission-text="Title is private" data-url="https://github.com/google/jax/issues/1733" data-hovercard-type="issue" data-hovercard-url="/google/jax/issues/1733/hovercard" href="https://github.com/google/jax/issues/1733">#1733</a></p>
<hr>
<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 requests.</li>
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> <a class="issue-link js-issue-link" data-error-text="Failed to load title" data-id="933465160" data-permission-text="Title is private" data-url="https://github.com/google/jax/issues/7145" data-hovercard-type="issue" data-hovercard-url="/google/jax/issues/7145/hovercard" href="https://github.com/google/jax/issues/7145">#7145</a></li>
</ul>
<p dir="auto">Running <code class="notranslate">import jax</code> with jaxlib 0.1.52 and jax 0.1.75 on mac osx 13.16 crashes and gives this error:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="python
Python 3.8.2 | packaged by conda-forge | (default, Mar 23 2020, 17:55:48)
[Clang 9.0.1 ] on darwin
Type &quot;help&quot;, &quot;copyright&quot;, &quot;credits&quot; or &quot;license&quot; for more information.
&gt;&gt;&gt; import jax
Traceback (most recent call last):
  File &quot;&lt;stdin&gt;&quot;, line 1, in &lt;module&gt;
  File &quot;/Users/Chris/miniconda3/envs/python3/lib/python3.8/site-packages/jax/__init__.py&quot;, line 22, in &lt;module&gt;
    from .api import (
  File &quot;/Users/Chris/miniconda3/envs/python3/lib/python3.8/site-packages/jax/api.py&quot;, line 38, in &lt;module&gt;
    from . import core
  File &quot;/Users/Chris/miniconda3/envs/python3/lib/python3.8/site-packages/jax/core.py&quot;, line 31, in &lt;module&gt;
    from . import dtypes
  File &quot;/Users/Chris/miniconda3/envs/python3/lib/python3.8/site-packages/jax/dtypes.py&quot;, line 31, in &lt;module&gt;
    from .lib import xla_client
  File &quot;/Users/Chris/miniconda3/envs/python3/lib/python3.8/site-packages/jax/lib/__init__.py&quot;, line 52, in &lt;module&gt;
    from jaxlib import xla_client
  File &quot;/Users/Chris/miniconda3/envs/python3/lib/python3.8/site-packages/jaxlib/xla_client.py&quot;, line 39, in &lt;module&gt;
    from . import xla_extension as _xla
ImportError: dlopen(/Users/Chris/miniconda3/envs/python3/lib/python3.8/site-packages/jaxlib/xla_extension.so, 2): Symbol not found: ____chkstk_darwin
  Referenced from: /Users/Chris/miniconda3/envs/python3/lib/python3.8/site-packages/jaxlib/xla_extension.so (which was built for Mac OS X 10.15)
  Expected in: /usr/lib/libSystem.B.dylib
 in /Users/Chris/miniconda3/envs/python3/lib/python3.8/site-packages/jaxlib/xla_extension.so
&gt;&gt;&gt;"><pre class="notranslate"><code class="notranslate">python
Python 3.8.2 | packaged by conda-forge | (default, Mar 23 2020, 17:55:48)
[Clang 9.0.1 ] on darwin
Type "help", "copyright", "credits" or "license" for more information.
&gt;&gt;&gt; import jax
Traceback (most recent call last):
  File "&lt;stdin&gt;", line 1, in &lt;module&gt;
  File "/Users/Chris/miniconda3/envs/python3/lib/python3.8/site-packages/jax/__init__.py", line 22, in &lt;module&gt;
    from .api import (
  File "/Users/Chris/miniconda3/envs/python3/lib/python3.8/site-packages/jax/api.py", line 38, in &lt;module&gt;
    from . import core
  File "/Users/Chris/miniconda3/envs/python3/lib/python3.8/site-packages/jax/core.py", line 31, in &lt;module&gt;
    from . import dtypes
  File "/Users/Chris/miniconda3/envs/python3/lib/python3.8/site-packages/jax/dtypes.py", line 31, in &lt;module&gt;
    from .lib import xla_client
  File "/Users/Chris/miniconda3/envs/python3/lib/python3.8/site-packages/jax/lib/__init__.py", line 52, in &lt;module&gt;
    from jaxlib import xla_client
  File "/Users/Chris/miniconda3/envs/python3/lib/python3.8/site-packages/jaxlib/xla_client.py", line 39, in &lt;module&gt;
    from . import xla_extension as _xla
ImportError: dlopen(/Users/Chris/miniconda3/envs/python3/lib/python3.8/site-packages/jaxlib/xla_extension.so, 2): Symbol not found: ____chkstk_darwin
  Referenced from: /Users/Chris/miniconda3/envs/python3/lib/python3.8/site-packages/jaxlib/xla_extension.so (which was built for Mac OS X 10.15)
  Expected in: /usr/lib/libSystem.B.dylib
 in /Users/Chris/miniconda3/envs/python3/lib/python3.8/site-packages/jaxlib/xla_extension.so
&gt;&gt;&gt;
</code></pre></div>
<p dir="auto">Jax 0.1.75 and jaxlib 0.1.51 together work fine, so this bug seems like it's been introduced in jaxlib 0.1.51.</p>