This is also a remanding of the original PR by @wdphy16 , following the
new infrastructure for Jax-compatible operators.

resuscitates #968

@wdphy16 as you are the original author of the code, could you please go
back through it and comment it throughout? It's very complicated and
it's very hard to follow what the generation of masks do...

---------

authored-by: Clemens Giuliani <clemens@inailuig.it>