How to avoid data copy working with numpy arrays and JAX on the CPU?

[r-aristov]

Hello. I can not figure out how to implement the following: I'd like to have a large buffer in RAM (on CPU) and write items to it without creating a copy, like I can do with numpy - buffer[ptr] = data. But I also want to sample batches from it fast and move them to GPU.

I thought I can use JITed JAX function that efficiently handles sampling, but ran into several problems. Even on the CPU JAX can not modify data in place - am I correct, or did I miss something?

But this is not a real issue, my real question is following. Trying to solve this, I've made a numpy array CPU buffer, which allows in-place modifications. And tried to convert it to JAX array on the CPU for sampling without copying the whole buffer. But JAX still creates a copy of the data, it does not matter whether I call jnp.asarray or jnp.array.
I've read several posts here with similar issue and used the suggested code to check if copying occures:

import numpy as np
import jax.numpy as jnp
import jax
import resource
jax.config.update('jax_platform_name', 'cpu')

rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss/1024
a = np.ones(int(16e8)) 
rss1 = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss/1024
print('np array allocated:', rss1 - rss)
b = jnp.asarray(a, copy=False)
rss2 = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss/1024
print('moved to JAX', rss2 - rss1)
print('JAX array device:', b.devices())

This code prints:

np array allocated: 12207.02734375
moved to JAX 6228.19140625
JAX array device: {CpuDevice(id=0)}

First of all, I can not quite understand where these numbers are coming from - I mean, why there is 6 GB increase and not 12 GB if this is a copy?
Second, if I use context manager, it seems, it does not set the device properly -

# use context manager instead of the global config setting
# jax.config.update('jax_platform_name', 'cpu')
# ...

with jax.default_device(jax.devices("cpu")[0]):
    b = jnp.asarray(a, copy=False)

I get this error:

ValueError: jnp.asarray: cannot convert object of type <class 'numpy.ndarray'> to JAX Array on backend='gpu' with copy=False. Consider using copy=None or copy=True instead.

Which suggests that it sees backend=gpu.
Can you help me to resolve this?

[jakevdp]

Your NumPy array is of type float64, and by default, JAX cannot use 64-bit values, so it is converted to a float32 array for use by JAX.

If you'd like to use 64-bit values in JAX, you could set the X64 flag to True (see https://docs.jax.dev/en/latest/default_dtypes.html)

Alternatively, if you'd like to continue using float32, values, you should allocate the NumPy array in float32:

a = np.ones(int(16e8), dtype='float32')

With this change, jnp.asarray should reuse the NumPy buffer as long as it is byte-aligned.

[r-aristov]

Ah, that explains why there is a 6 GB increase instead of 12 GB!
Thank you very much, it solved half of the issue.

But I still can not use the context manager to set the target device = CPU for this operation! This looks like a bug to me. Or am I missing something?

import numpy as np
import jax.numpy as jnp
import jax
import resource
# jax.config.update('jax_platform_name', 'cpu')

rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss/1024
a = np.ones(int(16e8), dtype='float32') 

rss1 = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss/1024
print('np array allocated:', rss1 - rss)
with jax.default_device(jax.devices("cpu")[0]):
    b = jnp.asarray(a, copy=False)
rss2 = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss/1024
print('moved to JAX', rss2 - rss1)
print('JAX array device:', b.devices())

Traceback (most recent call last):
File "", line 2, in
File ".../jax/_src/numpy/lax_numpy.py", line 3283, in asarray
raise ValueError(f"jnp.asarray: cannot convert object of type {type(a)} to JAX Array "
ValueError: jnp.asarray: cannot convert object of type <class 'numpy.ndarray'> to JAX Array on backend='gpu' with copy=False. Consider using copy=None or copy=True instead.

[r-aristov]

UPD: I see, there is a difference between default device and default backend.
This code seems to work as expected and does not copy the data if np array type is compatible.

b = jax.device_put(a, device=jax.devices("cpu")[0])

Well, that solves the issue. It's a bit unexpected that context manager does not affect a backend though.

I have one related question left - I am trying to remove unnecessary data copying from my pipeline.
In current implementation (with the issues above fixed) I use np array as a temporary buffer since it allows in-place operations.

But is there a way to efficiently sample directly from the list of arrays without combining them into one in JAX?
To avoid copying data into one large continuous numpy array and work directly with the list of JAX arrays (named tuples with JAX arrays as members, but it is irrelevant here) that is populated by a worker thread, that produces one array at a time and adds them to the list?
It seems that this is a right task for JAX - I mean, the size of the list is fixed (it is list(buffer_deque) with fixed max_len), samples from each of the arrays can be drawn in parallel.

One way to do it is to use python for loop to run through a list and sample items from each of the arrays - when JITed it is petty fast, but JITting takes ages even for a small buffer size of 100-200 items.
I tried to use JAX fori loop or scan, but it seem that it does not work that way - I can not index list from inside the JITed function.