Generating arrays of different shapes can be 1000 time slower than numpy, is there a better way to do it?

[morgangiraud]

Hi everyone, thanks for the work done on JAX, I've been quite pleased so far with the performance of the lib. Now I'm running into a problem for which I can't find a good way to do it.

I'm trying to create a high number of arrays for which their shape are generated randomly. I understand that JAX add overheads to this kind of small operation but is there a way to generate those arrays efficiently?

Here is the code I've been using:

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

sizes_np = np.hstack([
    np.array([1] * nb_arrays)[:, np.newaxis], 
    np.random.randint(2**3, 2**6, [nb_arrays, 2])
])

t0 = time.time()
for i in range(100):
    shape = sizes_np[i]
    mask_np = np.ones(shape)
t1 = time.time()

print(t1 - t0)


key = jax.random.PRNGKey(0)
sizes = jnp.hstack([
    jnp.array([1] * nb_arrays)[:, jnp.newaxis], 
    jax.random.randint(key, [nb_arrays, 2], 2**2, 2**3)
])

t0 = time.time()
for i in range(100):
    shape = sizes[i]
    mask = jnp.ones(shape)
t1 = time.time()

print(t1 - t0)

Outputs:

0.0002486705780029297
0.23677420616149902

Note: The end goal is to pad those arrays and then vstack them into one big array of ones and zeros (The important part is to keep those rectangles of ones)