tf.computeWeightedLoss should apply broadcast before reduction

[manrajgrover]

To get help from the community, check out our Google group.

TensorFlow.js version

0.11.9

Browser version

Version 67.0.3396.87 (Official Build) (64-bit)

Describe the problem or feature request

tf.computeWeightedLoss currently does not apply broadcast on weights on reduction SUM_BY_NONZERO_WEIGHTS before calculating non-zero values. This leads to wrong output. See tensorflow implementation. Found it while working on tensorflow/tfjs-core#1120

Code to reproduce the bug / link to feature request

const losses = tf.tensor2d([[0, 0, 1], [1, 0, 0], [0, 1, 0]], [3, 3]);
const weights = tf.tensor2d([[0.1, 0.2, 0.3]]);

tf.losses.computeWeightedLoss(losses, weights).print();
// Tensor
//    0.20000001788139343 (should be 0.06666667014360428)

[manrajgrover]

I can take this one up.