How to subset a large matrix of data in jax.lax.scan?

[kclauw]

Hi,

I have a large matrix of data of size MxN and a list of variable size indices (multiplets). The goal is to run a function over a subset of the data by slicing the matrix into smaller batches of size MxK where K is the length of the indices and K <= M . To illustrate this, I have made a toy example using the following code:

import jax
import jax.numpy as jnp


from functools import partial


data = jnp.array([np.arange(32) for i in range(10)]) #NxM data matrix (this is to large to pass in batches)

#Add padding to the multiplets to make a static array (this is just for illustration purposes)
padding = jnp.nan
multiplets = jnp.array([
    [0, 1, 2, padding, padding], 
    [5, 6, padding, padding, padding], 
    [9, 10, 11, 32, 24]
])

def process_data(x):
    #This is a more complicated function that will operate on a subset of the data on the columns M
    return x 

@partial(jax.jit, static_argnums=(2, 3))
def subset(carry, index):
    data, multiplets = carry
    subset_data = data[:, ~jnp.isnan(multiplets[index])] #How do I slice the multiplets and ignore the nan values? This part does not work
    results = process_data(subset_data)
    return carry, results

final, result = jax.lax.scan(subset, (data, multiplets), jnp.arange(len(multiplets)))

My approach is to:

1. Add padding with nan values to make it static
2. Use only the indices without nan values but this slicing does not seem to work as I get an error: This code receives the following error: ax.errors.NonConcreteBooleanIndexError: Array boolean indices must be concrete; got ShapedArray(bool[5].

Does anyone have any advice on how to approach this? Is there perhaps a better way? I found similar discussion using padding but - to the best of my knowledge - I have not seen anyone using the indices with padding to select items from a matrix. Thanks

[jakevdp]

Thanks for the question. The issue you're running into is that you cannot use dynamically-shaped arrays within jax.scan and other JAX transformations, and data[:, ~jnp.isnan(multiplets[index])] has a dynamic shape (i.e. a shape that depends on the values in a traced array). See JAX Sharp Bits: Dynamic Shapes for more information on this.

I suspect your best approach here will be to make subset_data itself a padded array. You could use out-of-bound indices for your padding value, and then handle the padded values in process_data. Something like this:

import jax
import jax.numpy as jnp


from functools import partial


data = jnp.array([jnp.arange(32.0) for i in range(10)]) #NxM data matrix (this is to large to pass in batches)


padding = data.shape[1]  # out-of-bound index
multiplets = jnp.array([
    [0, 1, 2, padding, padding], 
    [5, 6, padding, padding, padding], 
    [9, 10, 11, 32, 24]
])

def process_data(x):
    #This is a more complicated function that will operate on a subset of the data on the columns M
    return jnp.nansum(x)

@jax.jit
def subset(carry, index):
    data, multiplets = carry
    subset_data = data.at[:, multiplets[index]].get(mode='fill', fill_value=jnp.nan)
    results = process_data(subset_data)
    return carry, results

final, result = jax.lax.scan(subset, (data, multiplets), jnp.arange(len(multiplets)))
print(result)
# [ 30. 110. 540.]

Is that something like what you have in mind?

[kclauw]

This is what I had in mind. Thanks.