jnp.asarray making copies in latest jax

[mattbarrett98]

Description

import numpy as np
import jax.numpy as jnp
import os, psutil

process = psutil.Process(os.getpid())
print(process.memory_info().rss / 1e9)  # 0.1 GB
x = np.random.uniform(0, 1, (10000, 20000)).astype("float32")
print(process.memory_info().rss / 1e9)  # 0.9 GB
y = jnp.asarray(x)
print(process.memory_info().rss / 1e9)  # 1.7 GB

In version 0.4.16, jnp.asarray now seems to be making copies when called on numpy arrays. In earlier versions, it would just reuse the underlying data and no (little) extra memory would be consumed. Is this intended?

What jax/jaxlib version are you using?

jax, jaxlib 0.4.16

Which accelerator(s) are you using?

CPU

Additional system info

No response

NVIDIA GPU info

No response

[hawkinsp]

Yash, is it possible that this occurs because of the changes to the primitive dispatch path and the fact that we now copy for trivial computations?

[yashk2810]

Yeah, that sounds correct to me.

We removed the trivial dispatch path from JAX because it was a pre-ominstaging optimization which doesn't have much use post-omnistaging. Also it helped speed up the dispatch path by 100x.

Is the copy causing problems?

[hawkinsp]

I think not copying (while it has only ever been best effort) is an important optimization.

[yashk2810]

#17721 should fix

[mattbarrett98]

Thanks guys! before the next release is there another way of converting without copies?

[jakevdp]

If you want to convert a numpy array to a jax array without copying in version 0.4.16, you can use y = jax.device_put(x)

This will only be avoid copies if using a CPU device, and if your numpy array's byte alignment is compatible with XLA.