Unnecessary loss of precision when computing loss functions about dfdx HOT 2 CLOSED

This was done specifically for f16 support - notably the max value an f16 can store is 65504, and it has low precision for generally high values.

In the scalar version of div/mul the operands will be converted to the dtype before executing - so the f64 would be converted to an f16 before actually running. By using 1 / a, we have a better chance that f16 can store the actual value.

from dfdx.

We may also be able to replace the call to mean() with a call to sum(), but it requires fiddling with the tensor dimensions. Using mean() and rescaling by the last axis size calculates a sum of the form $$\text{mean} \cdot \text{last axis numel} = \left(\dfrac{1}{|A||I|}\sum_{(a, i)\in A\times I}p_{a, i}\right)\cdot |A|$$ whereas we could instead compute $$\text{sum} / \text{foo} = \left(\sum_{(a, i)\in A\times I}p_{a, i}\right)\cdot |I|^{-1}.$$

I couldn't immediately see how to compute for $|I|$ using the generic parameter S or its axis using the existing trait bounds.

from dfdx.