Masking in Einops

[JTrainer-STORi]

For the reduce function, specify a mask argument. The reduce function would then be applied to the entries in the array for which the mask argument = True. For an example of how this would work, take an example array array = [[1, 3, 4],[5, 4, 1],[8, 0, 6]] and mask = [True, False, True]. Calling reduce(array, 'a b -> b', 'sum', mask) would then return [9, 3, 10] ignoring the middle row of the array for which the mask = False in the summation.

This feature would be personally useful for me because I enjoy the syntax of einops and I'm currently in a situation where I need to implement masking in a neural network. Being able to use the above syntax would make the code as a whole much cleaner.

[arogozhnikov]

Hi @JTrainer-STORi , thx for suggestion. I don't think it is a good addition to reduce:

• your specific scenario is very simple to address with existing tools (basically, array[mask], then reduce).
• should einops deal with this, it implies ndim masking, and I'm not sure there is any consensus between frameworks that einops can rely on.
• it is not that easy to read from operation that mask should have shape a. One can assume that mask can have a shape a b and summation should happen only over unmasked elements. May result in silent problems for square matrices

[JTrainer-STORi]

Hi @arogozhnikov thanks for the quick response, indexing the array before passing it to reduce didn't cross my mind for some reason but as you say it does what I want.