The formula before was wrong, it does not work when inputs are matrices,
because it depends on which axis y is computed.  I believe this formulation
is much better.