scaning or similar over a list of functions

[pabloduque0]

Lets assume we have a list of functions (like data augmentation functions for example) in which all transformations are applied to the output of the previous function in a sequential way. Something like:

images = jnp.ones((25, 100, 100, 3))
augmentations = [fn1, fn2, fn3]

for aug in augmentations:
    images = aug(images)

scan seemed like a great choice since I can rely on the carry but after a couple attempts I think I cannot scan over a list of functions. vmap does not necessarily fit in here either and I think we cannot vmap over such list anyways.

I was wondering if there is a way to speed up this beyond jiting the function where this logic happens. Maybe there is actually not greater benefit from trying something else and jit could be as good as we are going to get.

Thank you very much in advance!

[jakevdp]

It is only possible to scan over a list of values, not over a list of functions. The best way to achieve what you have in mind is likely with a for loop, as you used in your example above.

[pharringtonp19]

@pabloduque0 I had a similar question that I asked to the folks leading the dex-lang project. You might find the comments there helpful as well.

[pabloduque0]

Thank you both for the answers, will take a look at that!