loop over unequal intervals

[nestordemeure]

I have a function that looks like this (the intervals have unequal sizes, indices is an array of integer indices):

def f(input, output, intervals):
   for start,end in intervals:
      input_interval = input[indices, start:end, :]
      output_interval = output[indices, start:end]
      output[indices, start:end] = kernel(input_interval, output_interval)
   return output

What would be the recommended pattern to write it in JAX? At the moment I am only jitting kernel but I feel that jitting the full function could reduce memory transfers.