Inserting values to jax arrays

[yardenas]

I'm learning jax by implementing my own version of Dreamer (see here). To be able to jit to whole training loop (sample batch -> update model, actor, critic -> sample batch ->....), I implemented the replay buffer with jax arrays. So, for instance, given an image observation of shape (64, 64, 3), a pre-defined episode length of 1000 steps and a capacity of (say) 50 episodes, I allocate an array in the following way:

{
'observation': jax.device_put(jnp.full((capacity, max_episode_length + 1) + (64, 64, 3), jnp.nan, jnp.uint8), device),
...
}

where device points to my CPU.
When new data arrives, I insert it to the buffer:

...
self.data['observation'] = self.data['observation'].at[self.idx, position].set(data)
...

Unfortunately, this seems like a very slow operation, and what's more weird for me is that the bigger the capacity is, the slower this operation becomes.

I've created two gists that demonstrate this behavior:

1. A script that does the procedure explained above, taking as a parameter a capacity. So for example, on my (CPU) system, running python3 insert_slow_jax.py --capacity 50 yields an average insertion time of ~0.4 seconds while running python3 insert_slow_jax.py --capacity 2 yields an insertion time of ~0.02 seconds.
2. Another script that (hopefully) does exactly the same, but implemented in numpy, yielding similar and much faster insertion times with both capacities (2 and 50).

I created the second gist not because I wanted to compare insertion times between jax and numpy (this is already explained nicely on the FAQ), but mostly because I wasn't sure if the numpy script would also run slower when increasing the buffer's capacity.

Any thoughts about this?
I might probably relax the constraint of jiting the whole training loop and just sample numpy arrays from the replay buffer, but I'm still interested on what is the cause of this behavior.

[jakevdp]

Outside JIT, insertion operations like x = x.at[y].set(z) result in a copy of the full buffer. So it makes sense that (1) it would be slower than numpy's in-place insertions, and (2) it would be slower for larger buffers than for smaller buffers.

Inside JIT, XLA is able to avoid these copies by essentially compiling them into in-place operations.

So I'd suggest JIT-compiling your function to improve the performance.

[yardenas]

Thank you! It does make a lot of sense to me now.
I continued playing around with this small example and wrapped the store(...) function with jit:

import argparse
import functools
import timeit
from typing import Mapping, Union

import jax
import jax.numpy as jnp
import numpy as np


class ReplayBuffer:
    def __init__(
            self,
            capacity: int,
            max_episode_length: int,
            observation_space,
            batch_size: int,
            length: int
    ):
        device = jax.devices("cpu")[0]
        observation = jax.device_put(jnp.full(
            (capacity, max_episode_length + 1) + observation_space.shape,
            jnp.nan, jnp.uint8), device)
        self.data = {
            'observation': observation
        }
        self.episdoe_lengths = jnp.full((capacity,), 0, dtype=jnp.uint32)
        self.idx = 0
        self.capacity = capacity

    @functools.partial(jax.jit, static_argnums=0)
    def store(self, transition, buffer, idx, lengths):
        position = lengths[idx]
        for key in self.data.keys():
            buffer[key] = buffer[key].at[idx, position].set(transition[key])
        lengths = lengths.at[idx].add(1)
        return buffer, lengths, idx


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--capacity', default=50, type=int)
    config = parser.parse_args()
    rb = ReplayBuffer(config.capacity, 1000, np.ones((64, 64, 3)), 32, 50)

    def insert_once():
        dummy_o = np.random.randint(0, 255, (64, 64, 3)).astype(np.float32)
        rb.store(dict(observation=dummy_o), rb.data, rb.idx, rb.episdoe_lengths)

    print("Jitting...")
    insert_once()
    print("Done jitting...")
    t = timeit.Timer(insert_once)
    print("Measuring insertion time...")
    print("Done measuring. Average time:", t.timeit(number=100) / 100)


if __name__ == '__main__':
    main()

However, the only way I could increase performance was by commenting out the return statement of store. I guess that another copy occurs when returning the buffer from store; is that correct?

Thank you again!

[jakevdp]

Yes, in general buffers will be copied at JIT boundaries (i.e. the buffer you get out of a function call will be a copy of the input buffer) but if you have a larger program within the JIT compilation, invocations of this function will not lead to buffer copies, because XLA can compile them away.

A couple other notes:

• you should be careful about Asynchronous dispatch when writing benchmarks like this.
• your insert_once function is creating a large numpy array and then pushing it to the device; this device transfer cost is often not negligible, so you might want to separate that out from your benchmark if you're mainly interested in the speed of insertions.

You might read through Benchmarking JAX Code to make sure you're measuring what you think you're measuring.

[yardenas]

Thank you for the detailed response.