Issue 1:
I wanted to have two pmap()-ed function, piping one's output into another. For example, in the following way:

for i, epoch in enumerate(tqdm(range(1, num_train_epochs + 1), desc=f"Epoch ...", 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()

parallel_train_step() and parallel_eval_step() are the two pmap()-ed functions and state is the object being passed from one into another. The above code is adapted from this HuggingFace Notebook.

Now, in case of small models, the state is small in size and the output of parallel_train_step or parallel_eval_step gets stored in a buffer distinct from the input buffer. But in case of large models, the output state may get stored in the input buffer itself. But the input buffer is marked as donated, so when the output state is passed into the second function(parallel_eval_step()), it raises RuntimeError: Invalid argument: Buffer has been deleted or donated .

I notice that we can print() the returned state, we can also serialize() it, but can't pass the buffers to the second pmap()-ed function.

Manually, I tried solving the above error by using state = jax.tree_util.tree_map(lambda x: jnp.array(x.tolist()), state) before passing it to the second function. Also, tried using lambda x: jnp.array(x,copy=True). Both of which didn't work. Is there some other way to solve the problem?

Possible solution can be to un-mark the buffer as donated when the input buffer is used to store the output. But how to do it?

Related: #1733

Please:

Check for duplicate requests.

#7145

Issue 2:
Running import jax with jaxlib 0.1.52 and jax 0.1.75 on mac osx 13.16 crashes and gives this error:

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.
>>> import jax
Traceback (most recent call last):
File "<stdin>", line 1, in <module>
File "/Users/Chris/miniconda3/envs/python3/lib/python3.8/site-packages/jax/__init__.py", line 22, in <module>
from .api import (
File "/Users/Chris/miniconda3/envs/python3/lib/python3.8/site-packages/jax/api.py", line 38, in <module>
from . import core
File "/Users/Chris/miniconda3/envs/python3/lib/python3.8/site-packages/jax/core.py", line 31, in <module>
from . import dtypes
File "/Users/Chris/miniconda3/envs/python3/lib/python3.8/site-packages/jax/dtypes.py", line 31, in <module>
from .lib import xla_client
File "/Users/Chris/miniconda3/envs/python3/lib/python3.8/site-packages/jax/lib/__init__.py", line 52, in <module>
from jaxlib import xla_client
File "/Users/Chris/miniconda3/envs/python3/lib/python3.8/site-packages/jaxlib/xla_client.py", line 39, in <module>
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
>>>

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.