Vmap of Python Lists and JAX StringArray

[GoktugGuvercin]

Hello JAX Community;

To vectorize data loading process in a typical data-loader, I wanted to use jax.vmap. The names of images are stored in a list, one image name is fetched from the list, and that name is used with PIL.Image.open() function to load corresponding image. You can see the code below.

def load(self, index):
    img_dir = os.path.join(self.root_dir, self.img_names[index])
    img = open(img_dir, mode="r", formats=None)
    return img

vmap_load = jax.vmap(self.load, in_axes=0, out_axes=0)

When I try to vectorize it, I faced TracerIntegerConversionError since probably what I try to do it is actually indexing a python list with a tracer object. To solve this problem, I wanted to convert the list of file names into JAX Array, but then I noticed that StringArray is not supported in JAX. The solution coming to the mind is to change self.img_names into classic numpy array to do multiple indexing, but even if I do that, I think that it is not possible to vectorize open() function without jax.vmap.

Is there any other solution to vectorize load() function that you can recommend ?
Is there any other way to support strings in JAX arrays ?

Thanks in advance.

[jakevdp]

There is not really any way to do what you have in mind – JAX doesn't support operations on strings. There are mechanisms you can use to refer to arrays of string data using JAX constructs (see for example the StringArray implementation in #12329) but they won't help you in your goal of vectorizing file loading, because file loading is not limited by your ability to call load() in parallel, but rather by your computers ability to physically read data from the disk in parallel.

[GoktugGuvercin]

Oh, I did not think about this in terms of hardware parallelism. Yes, you are totally right. Thank you so much for your explanation.