docs: @Override
      public INDArray computeGradient(INDArray labels, INDArray preOutput, IActivation activationFn, INDArray mask) {
          labels = labels.castTo(preOutput.dataType());   //No-op if already correct dtype
          long nSamples = labels.size(0);
  
          INDArray output = activationFn.getActivation(preOutput.dup(), false);
  
          MixtureDensityComponents mdc = extractComponents(output);
  
          INDArray gradient = Nd4j.zeros(nSamples, preOutput.columns());
  
          INDArray labelsMinusMu = labelsMinusMu(labels, mdc.mu);
          INDArray labelsMinusMuSquared = labelsMinusMu.mul(labelsMinusMu).sum(2);
  
          // This computes pi_i, see Bishop equation (30).
          // See http://www.plsyard.com/dealing-overflow-and-underflow-in-softmax-function/
          // this post for why we calculate the pi_i in this way.
          // With the exponential function here, we have to be very careful
          // about overflow/underflow considerations even with
          // fairly intermediate values.  Subtracting the max
          // here helps to ensure over/underflow does not happen here.
          // This isn't exactly a softmax because there's an 'alpha' coefficient
          // here, but the technique works, nonetheless.
          INDArray variance = mdc.sigma.mul(mdc.sigma);
          INDArray minustwovariance = variance.mul(2).negi();
          INDArray normalPart = mdc.alpha.div(Transforms.pow(mdc.sigma.mul(SQRT_TWO_PI), mLabelWidth));
          INDArray exponent = labelsMinusMuSquared.div(minustwovariance);
          INDArray exponentMax = exponent.max(1);
          exponent.subiColumnVector(exponentMax);
          INDArray pi = Transforms.exp(exponent).muli(normalPart);
          INDArray piDivisor = pi.sum(true,1);
          pi.diviColumnVector(piDivisor);
  
          // See Bishop equation (35)
          //INDArray dLdZAlpha = Nd4j.zeros(nSamples, nLabelsPerSample, mMixturesPerLabel); //mdc.alpha.sub(pi);
          INDArray dLdZAlpha = mdc.alpha.sub(pi);
          // See Bishop equation (38)
          INDArray dLdZSigma = (labelsMinusMuSquared.div(variance).subi(mLabelWidth)).muli(-1).muli(pi);
          // See Bishop equation (39)
  
          // This turned out to be way less efficient than
          // the simple 'for' loop here.
          //INDArray dLdZMu = pi
          //        .div(variance)
          //        .reshape(nSamples, mMixtures, 1)
          //        .repeat(2, mLabelWidth)
          //        .muli(labelsMinusMu)
          //        .negi()
          //        .reshape(nSamples, mMixtures * mLabelWidth);
  
          INDArray dLdZMu = Nd4j.create(nSamples, mMixtures, mLabelWidth);
          for (int k = 0; k < mLabelWidth; k++) {
              dLdZMu.put(new INDArrayIndex[] {NDArrayIndex.all(), NDArrayIndex.all(), NDArrayIndex.point(k)},
                              labelsMinusMu.get(new INDArrayIndex[] {NDArrayIndex.all(), NDArrayIndex.all(),
                                              NDArrayIndex.point(k)}).muli(pi).divi(variance).negi());
          }
          dLdZMu = dLdZMu.reshape(nSamples, mMixtures * mLabelWidth);
  
          // Place components of gradient into gradient holder.
          gradient.put(new INDArrayIndex[] {NDArrayIndex.all(), NDArrayIndex.interval(0, mMixtures)}, dLdZAlpha);
          gradient.put(new INDArrayIndex[] {NDArrayIndex.all(), NDArrayIndex.interval(mMixtures, mMixtures * 2)},
                          dLdZSigma);
          gradient.put(new INDArrayIndex[] {NDArrayIndex.all(),
                          NDArrayIndex.interval(mMixtures * 2, (mLabelWidth + 2) * mMixtures)}, dLdZMu);
  
          INDArray gradients = activationFn.backprop(preOutput, gradient).getFirst();
  
          if (mask != null) {
              LossUtil.applyMask(gradients, mask);
          }
  
          return gradients;
      }