Best way compute dynamic slices of arrays where the slices are known beforehand?

[Mattias421]

I am trying to align sequences of text feature vectors with sequences of speech feature vectors.
My data is organised with text_feature_array which is an array of text feature vectors where each vector represents a character/token in the dataset. You can think of this as flattening out the entire dataset and then extracting feature vectors. The length of each text segment is recorded in a separate array text_lens, which allows the text features for each sample to be extracted using cumulative sum and text_lens, e.g. the 0th text feature segment would be text_feature_array[0 : 0 + textlens[0]], and the 1th segment would be text_feature_array[0 + textlens[0]: 0 + textlens[0]] + textlens[1]] and so on. The same applies to speech_feature_array, but the lengths are much larger than the text segments.

I have a function align that, given a text and speech segment, will assign a duration/number to each text feature vector in the segment that assigns how many speech vectors correspond to it, align returns a vector with length equal to the text segment and the sum of align is equal to the length of the speech segment.

I'm trying to make use of Jax's flexible concurrency to compute all the alignments as fast as possible, I have considered use jax.lax.scan like so (I've simplified the code for readability):

def f(durations, indicies):
  txt_idx, txt_len, sp_idx, sp_len = indicies
  text_segment = text_feature_array[txt_idx:txt_idx + txt_len]
  speech_segment = speech_feature_array[sp_idx:sp_idx + sp_len]
  durations[txt_idx:txt_idx + txt_len] = align(text_segment, speech_segment)
  return durations, 0

index_tuple = (text_cum_lens, text_lens, speech_cum_lens, speech_lens)
durations = zeros_like(text_feature_vector)
durations = jax.lax.scan(f, durations, xs=index_tuple)

My main issue is that txt_idx, txt_len, sp_idx, sp_len become abstract tracers meaning they cannot be used to slice the feature arrays. I have tried lax.dynamic_slice_in_dim but the same problem persists.

As this problem involves processing two large data structures to produce a new large data structure, parallel computing is highly desirable/possible, but I cannot quite figure out what angle to go about this.

Any thoughts?

[jakevdp]

If you want those features to be static, you'll need to keep them static in your scan, which you could do by making them static integers and closing over them in the scan body function.

That said, it looks like you are scanning over these values, so by definition they cannot be static/constant from iteration to iteration.

To make the semantics more clear and allow me to better answer your question, could you provide a minimal reproducible example of what exactly you're hoping to do?

[Mattias421]

Hi Jake, thanks for your reply!

I have a minimal example as follows:

def test_process_data_min_example():
    import jax
    import jax.numpy as jnp

    n_data = 10
    n_feat = 4
    key = jax.random.PRNGKey(0)
    data_array = jax.random.normal(key, shape=(n_data, n_feat))

    lengths = jnp.arange(n_data) + 1 # arbitrary lengths irl
    data_idx = jnp.cumsum(lengths)
    data_idx = jnp.insert(data_idx, 0, 0)

    data_array = jnp.repeat(data_array, lengths, axis=0)

    new_data_array = jnp.zeros_like(data_array)

    def dummy_fn(data_sample):
        return data_sample

    for i in range(lengths.shape[0]):
        idx = data_idx[i]
        length = lengths[i]

        data_sample = data_array.at[idx:idx + length].get()

        new_sample = dummy_fn(data_sample)

        new_data_array.at[idx: idx + length].set(new_sample)

My goal is to write a function that processes an array like so, it is very large (about 10 GB) so I would like to process every data sample in parallel to reduce computation time. As stated earlier I've tried with some methods such as lax.scan but as you say the lengths and indexes change from iteration to iteration so some dynamic slicing is required that I've not figured out how to do yet XD

[jakevdp]

If you need dummy_fn to take an array of size length, then I think this is the best approach: since the compiler does not currently support dynamically-shaped expressions, you need to send a different set of instructions for each shape, and the for loop is a convenient way to do that.

That approach will generally lead to long compile times, though, so you might be able to improve things by adjusting how the loop is implemented. For example, if you could write your loop body like this:

    for i in range(lengths.shape[0]):
        idx = data_idx[i]
        length = lengths[i]
        new_data_array = dummy_fn(data_sample, idx, length, new_data_array)

where dummy_fn operates on the appropriate subset of data_sample and writes it into new_data_array, then so long as dummy_fn does not use dynamic shapes to accomplish this, you could convert this for loop into a scan without issue.

It's hard to say more without diving into what dummy_fn actually does, but what do you think?

[Mattias421]

Ah this looks like a good direction. Since jax compiles a new function for each input shape I'm worried that my current mindset will always lead to long compile times as each datasample read from data_array has a different length. So I'm also thinking padding every sample to 1024 to reduce compilation time (lengths range from 95 to 896). For example the function will take a padded sample, index, length, and new_array and return a new updated array. As the size of the inputs and outputs remains constant in this case I think it will have reasonable performance. This assumes that I won't run into memory faults so let's see 😄