text: /**
   * @license
   * Copyright 2020 Google LLC. All Rights Reserved.
   * Licensed under the Apache License, Version 2.0 (the "License");
   * you may not use this file except in compliance with the License.
   * You may obtain a copy of the License at
   *
   * http://www.apache.org/licenses/LICENSE-2.0
   *
   * Unless required by applicable law or agreed to in writing, software
   * distributed under the License is distributed on an "AS IS" BASIS,
   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   * See the License for the specific language governing permissions and
   * limitations under the License.
   * =============================================================================
   */
  import { customGrad } from '../../gradients';
  import { convertToTensor } from '../../tensor_util_env';
  import { assertShapesMatch } from '../../util';
  import { add } from '../add';
  import { expandShapeToKeepDim } from '../axis_util';
  import { cast } from '../cast';
  import { div } from '../div';
  import { exp } from '../exp';
  import { logSumExp } from '../log_sum_exp';
  import { Reduction } from '../loss_ops_utils';
  import { mul } from '../mul';
  import { neg } from '../neg';
  import { op } from '../operation';
  import { reshape } from '../reshape';
  import { scalar } from '../scalar';
  import { sub } from '../sub';
  import { sum } from '../sum';
  import { computeWeightedLoss } from './compute_weighted_loss';
  /**
   * Computes softmax cross entropy between logits and labels.
   *
   * Measures the probability error in discrete classification tasks in which
   * the classes are mutually exclusive (each entry is in exactly one class).
   * For example, each CIFAR-10 image is labeled with one and only one label: an
   * image can be a dog or a truck, but not both.
   *
   * `NOTE`: While the classes are mutually exclusive, their probabilities need
   * not be. All that is required is that each row of labels is a valid
   * probability distribution. If they are not, the computation of the gradient
   * will be incorrect.
   *
   * `WARNING`: This op expects unscaled logits, since it performs a softmax on
   * logits internally for efficiency. Do not call this op with the output of
   * softmax, as it will produce incorrect results.
   *
   * logits and labels must have the same shape, e.g. [batch_size, num_classes]
   * and the same dtype.
   * @param labels The labels array.
   * @param logits The logits array.
   * @param dim The dimension softmax would be performed on. Defaults to `-1`
   *     which indicates the last dimension.
   */
  function softmaxCrossEntropyWithLogits_(labels, logits, dim = -1) {
      if (dim === -1) {
          dim = logits.rank - 1;
      }
      if (dim !== logits.rank - 1) {
          throw Error(`Softmax cross entropy along a non-last dimension is not yet ` +
              `supported. Labels / logits was rank ${logits.rank} ` +
              `and dim was ${dim}`);
      }
      // Use a custom gradient for numerical stability.
      const customOp = customGrad((labels, logits, save) => {
          // Reference:
          //   1. http://cs231n.github.io/linear-classify/#softmax
          //   2. https://blog.feedly.com/tricks-of-the-trade-logsumexp/
          const keepDims = true;
          const lse = logSumExp(logits, [dim], keepDims);
          const logResult = sub(cast(logits, 'float32'), lse);
          save([labels, logResult]);
          const costVector = neg(mul(logResult, labels));
          const value = sum(costVector, [dim]);
          const gradFunc = (dy, saved) => {
              const [labels, logResult] = saved;
              const dyShape = expandShapeToKeepDim(dy.shape, [dim]);
              return [
                  mul(reshape(dy, dyShape), sub(cast(labels, 'float32'), exp(logResult))),
                  mul(reshape(dy, dyShape), sub(exp(logResult), cast(labels, 'float32'))),
              ];
          };
          return { value, gradFunc };
      });
      return customOp(labels, logits);
  }
  /**
   * Computes the softmax cross entropy loss between two tensors.
   *
   * If labelSmoothing is nonzero, smooth the labels towards 1/2:
   *
   *   newOnehotLabels = onehotLabels * (1 - labelSmoothing)
   *                         + labelSmoothing / numClasses
   *
   * @param onehotLabels One hot encoded labels
   *    [batch_size, num_classes], same dimensions as 'predictions'.
   * @param logits The predicted outputs.
   * @param weights Tensor whose rank is either 0, or 1, and must be
   *    broadcastable to `loss`  of shape [batch_size]
   * @param labelSmoothing If greater than 0, then smooth the labels.
   * @param reduction Type of reduction to apply to loss. Should be of type
   *    `Reduction`
   *
   * @doc { heading: 'Training', subheading: 'Losses', namespace: 'losses' }
   */
  function softmaxCrossEntropy_(onehotLabels, logits, weights, labelSmoothing = 0, reduction = Reduction.SUM_BY_NONZERO_WEIGHTS) {
      let $onehotLabels = convertToTensor(onehotLabels, 'onehotLabels', 'softmaxCrossEntropy');
      const $logits = convertToTensor(logits, 'logits', 'softmaxCrossEntropy');
      let $weights = null;
      if (weights != null) {
          $weights = convertToTensor(weights, 'weights', 'softmaxCrossEntropy');
      }
      assertShapesMatch($onehotLabels.shape, $logits.shape, 'Error in softmaxCrossEntropy: ');
      if (labelSmoothing > 0) {
          const labelSmoothingScalar = scalar(labelSmoothing);
          const one = scalar(1);
          const numClasses = scalar($onehotLabels.shape[1]);
          $onehotLabels =
              add(mul($onehotLabels, sub(one, labelSmoothingScalar)), div(labelSmoothingScalar, numClasses));
      }
      const losses = softmaxCrossEntropyWithLogits_($onehotLabels, $logits);
      return computeWeightedLoss(losses, $weights, reduction);
  }
  export const softmaxCrossEntropy = op({ softmaxCrossEntropy_ });
  //# sourceMappingURL=softmax_cross_entropy.js.map