text: "use strict";
  /**
   * @license
   * Copyright 2018 Google Inc. 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.
   * =============================================================================
   */
  Object.defineProperty(exports, "__esModule", { value: true });
  var engine_1 = require("../engine");
  var gradients_1 = require("../gradients");
  var tensor_util_env_1 = require("../tensor_util_env");
  var util = require("../util");
  var axis_util = require("./axis_util");
  var operation_1 = require("./operation");
  var tensor_ops_1 = require("./tensor_ops");
  /**
   * Computes the log(sum(exp(elements across the reduction dimensions)).
   *
   * Reduces the input along the dimensions given in `axis`. Unless `keepDims`
   * is true, the rank of the array is reduced by 1 for each entry in `axis`.
   * If `keepDims` is true, the reduced dimensions are retained with length 1.
   * If `axis` has no entries, all dimensions are reduced, and an array with a
   * single element is returned.
   *
   * ```js
   * const x = tf.tensor1d([1, 2, 3]);
   *
   * x.logSumExp().print();  // or tf.logSumExp(x)
   * ```
   *
   * ```js
   * const x = tf.tensor2d([1, 2, 3, 4], [2, 2]);
   *
   * const axis = 1;
   * x.logSumExp(axis).print();  // or tf.logSumExp(a, axis)
   * ```
   * @param x The input tensor.
   * @param axis The dimension(s) to reduce. If null (the default),
   *     reduces all dimensions.
   * @param keepDims If true, retains reduced dimensions with length
   *     of 1. Defaults to false.
   */
  /** @doc {heading: 'Operations', subheading: 'Reduction'} */
  function logSumExp_(x, axis, keepDims) {
      if (axis === void 0) { axis = null; }
      if (keepDims === void 0) { keepDims = false; }
      var $x = tensor_util_env_1.convertToTensor(x, 'x', 'logSumExp');
      var axes = util.parseAxisParam(axis, $x.shape);
      var xMax = $x.max(axes, true /* keepDims */);
      var a = $x.sub(xMax);
      var b = a.exp();
      var c = b.sum(axes);
      var d = c.log();
      var res = xMax.reshape(d.shape).add(d);
      if (keepDims) {
          var newShape = axis_util.expandShapeToKeepDim(res.shape, axes);
          return res.reshape(newShape);
      }
      return res;
  }
  /**
   * Computes the sum of elements across dimensions of a `tf.Tensor`.
   *
   * Reduces the input along the dimensions given in `axes`. Unless `keepDims`
   * is true, the rank of the `tf.Tensor` is reduced by 1 for each entry in
   * `axes`. If `keepDims` is true, the reduced dimensions are retained with
   * length 1. If axes has no entries, all dimensions are reduced, and a
   * `tf.Tensor` with a single element is returned.
   *
   * ```js
   * const x = tf.tensor1d([1, 2, 3]);
   *
   * x.sum().print();  // or tf.sum(x)
   * ```
   *
   * ```js
   * const x = tf.tensor2d([1, 2, 3, 4], [2, 2]);
   *
   * const axis = 1;
   * x.sum(axis).print();  // or tf.sum(x, axis)
   * ```
   *
   * @param x The input tensor to compute the sum over. If the dtype is `bool`
   *   it will be converted to `int32` and the output dtype will be `int32`.
   * @param axis The dimension(s) to reduce. By default it reduces
   *     all dimensions.
   * @param keepDims If true, retains reduced dimensions with size 1.
   */
  /** @doc {heading: 'Operations', subheading: 'Reduction'} */
  function sum_(x, axis, keepDims) {
      if (axis === void 0) { axis = null; }
      if (keepDims === void 0) { keepDims = false; }
      var $x = tensor_util_env_1.convertToTensor(x, 'x', 'sum');
      if ($x.dtype === 'bool') {
          $x = $x.toInt();
      }
      var axes = util.parseAxisParam(axis, $x.shape);
      // Use a custom gradient to bypass 2 gradient backprops since sum is used
      // extremely often.
      var customOp = gradients_1.customGrad(function (x) {
          var permutation = axis_util.getAxesPermutation(axes, x.rank);
          var reductionAxes = axes;
          var permutedX = x;
          if (permutation != null) {
              permutedX = x.transpose(permutation);
              reductionAxes = axis_util.getInnerMostAxes(reductionAxes.length, x.rank);
          }
          var gradFunc = function (dy) {
              var expandedDyShape = x.shape.slice();
              axes.forEach(function (axis) {
                  expandedDyShape[axis] = 1;
              });
              var expandedDy = dy.reshape(expandedDyShape);
              var derX = expandedDy.mul(tensor_ops_1.ones(x.shape, 'float32'));
              return derX;
          };
          var gradInputs = function (dy) {
              return { x: function () { return gradFunc(dy); } };
          };
          var attrs = { axes: reductionAxes };
          var value = engine_1.ENGINE.runKernelFunc(function (backend) { return backend.sum(permutedX, reductionAxes); }, { x: permutedX }, gradInputs, 'Sum', attrs);
          if (keepDims) {
              var newShape = axis_util.expandShapeToKeepDim(value.shape, axes);
              value = value.reshape(newShape);
          }
          return { value: value, gradFunc: gradFunc };
      });
      return customOp($x);
  }
  /**
   * Computes the product of elements across dimensions of a `tf.Tensor`.
   *
   * Reduces the input along the dimensions given in `axes`. Unless `keepDims`
   * is true, the rank of the `tf.Tensor` is reduced by 1 for each entry in
   * `axes`. If `keepDims` is true, the reduced dimensions are retained with
   * length 1. If `axes` has no entries, all dimensions are reduced, and a
   * `tf.Tensor` with a single element is returned.
   *
   * ```js
   * const x = tf.tensor1d([1, 2, 3]);
   *
   * x.prod().print();  // or tf.prod(x)
   * ```
   *
   * ```js
   * const x = tf.tensor2d([1, 2, 3, 4], [2, 2]);
   *
   * const axis = 1;
   * x.prod(axis).print();  // or tf.prod(x, axis)
   * ```
   *
   * @param x The input tensor to compute the product over. If the dtype is `bool`
   *   it will be converted to `int32` and the output dtype will be `int32`.
   * @param axis The dimension(s) to reduce. By default it reduces
   *     all dimensions.
   * @param keepDims If true, retains reduced dimensions with size 1.
   */
  /** @doc {heading: 'Operations', subheading: 'Reduction'} */
  function prod_(x, axis, keepDims) {
      if (axis === void 0) { axis = null; }
      if (keepDims === void 0) { keepDims = false; }
      var $x = tensor_util_env_1.convertToTensor(x, 'x', 'prod');
      if ($x.dtype === 'bool') {
          $x = $x.toInt();
      }
      var axes = util.parseAxisParam(axis, $x.shape);
      var permutation = axis_util.getAxesPermutation(axes, $x.rank);
      var reductionAxes = axes;
      var permutedX = $x;
      if (permutation != null) {
          permutedX = $x.transpose(permutation);
          reductionAxes = axis_util.getInnerMostAxes(reductionAxes.length, $x.rank);
      }
      var value = engine_1.ENGINE.runKernelFunc(function (backend) { return backend.prod(permutedX, reductionAxes); }, { permutedX: permutedX });
      if (keepDims) {
          var newShape = axis_util.expandShapeToKeepDim(value.shape, axes);
          value = value.reshape(newShape);
      }
      return value;
  }
  /**
   * Computes the mean of elements across dimensions of a `tf.Tensor`.
   *
   * Reduces `x` along the dimensions given in `axis`. Unless `keepDims` is
   * true, the rank of the `tf.Tensor` is reduced by 1 for each entry in `axis`.
   * If `keepDims` is true, the reduced dimensions are retained with length 1.
   * If `axis` has no entries, all dimensions are reduced, and a `tf.Tensor` with
   * a single element is returned.
   *
   * ```js
   * const x = tf.tensor1d([1, 2, 3]);
   *
   * x.mean().print();  // or tf.mean(a)
   * ```
   *
   * ```js
   * const x = tf.tensor2d([1, 2, 3, 4], [2, 2]);
   *
   * const axis = 1;
   * x.mean(axis).print();  // or tf.mean(x, axis)
   * ```
   *
   * @param x The input tensor.
   * @param axis The dimension(s) to reduce. By default it reduces
   *     all dimensions.
   * @param keepDims If true, retains reduced dimensions with size 1.
   */
  /** @doc {heading: 'Operations', subheading: 'Reduction'} */
  function mean_(x, axis, keepDims) {
      if (axis === void 0) { axis = null; }
      if (keepDims === void 0) { keepDims = false; }
      var $x = tensor_util_env_1.convertToTensor(x, 'x', 'mean');
      var axes = util.parseAxisParam(axis, $x.shape);
      var shapes = axis_util.computeOutAndReduceShapes($x.shape, axes);
      var reduceShape = shapes[1];
      var reduceSize = util.sizeFromShape(reduceShape);
      // Use a custom gradient to bypass 2 gradient backprops since mean is used
      // extremely often.
      var customOp = gradients_1.customGrad(function (x) {
          var reduceSizeScalar = tensor_ops_1.scalar(reduceSize);
          // Cast if needed.
          var xReduce = reduceSizeScalar.dtype === x.dtype ? x : x.cast(reduceSizeScalar.dtype);
          var res = xReduce.div(reduceSizeScalar);
          var value = res.sum(axis, keepDims);
          var gradFunc = function (dy) {
              var expandedDyShape = x.shape.slice();
              axes.forEach(function (axis) {
                  expandedDyShape[axis] = 1;
              });
              var expandedDy = dy.reshape(expandedDyShape);
              var derX = expandedDy.mul(tensor_ops_1.ones(x.shape, 'float32')).div(reduceSize);
              return derX;
          };
          return { value: value, gradFunc: gradFunc };
      });
      return customOp($x);
  }
  /**
   * Gradient helper function for the min and max operations.
   */
  function gradForMinAndMax(dy, y, xOrig, origAxes, permutedAxes) {
      if (y.rank < xOrig.rank) {
          y = y.reshape(axis_util.expandShapeToKeepDim(y.shape, origAxes));
      }
      if (dy.rank < xOrig.rank) {
          dy = dy.reshape(axis_util.expandShapeToKeepDim(dy.shape, origAxes));
      }
      return {
          x: function () {
              var dx = dy.mul(xOrig.equal(y).cast(dy.dtype));
              return permutedAxes == null ? dx : dx.transpose(permutedAxes);
          }
      };
  }
  /**
   * Computes the minimum value from the input.
   *
   * Reduces the input along the dimensions given in `axes`. Unless `keepDims`
   * is true, the rank of the array is reduced by 1 for each entry in `axes`.
   * If `keepDims` is true, the reduced dimensions are retained with length 1.
   * If `axes` has no entries, all dimensions are reduced, and an array with a
   * single element is returned.
   *
   * ```js
   * const x = tf.tensor1d([1, 2, 3]);
   *
   * x.min().print();  // or tf.min(x)
   * ```
   *
   * ```js
   * const x = tf.tensor2d([1, 2, 3, 4], [2, 2]);
   *
   * const axis = 1;
   * x.min(axis).print();  // or tf.min(x, axis)
   * ```
   *
   * @param x The input Tensor.
   * @param axis The dimension(s) to reduce. By default it reduces
   *     all dimensions.
   * @param keepDims If true, retains reduced dimensions with size 1.
   */
  /** @doc {heading: 'Operations', subheading: 'Reduction'} */
  function min_(x, axis, keepDims) {
      if (axis === void 0) { axis = null; }
      if (keepDims === void 0) { keepDims = false; }
      var $x = tensor_util_env_1.convertToTensor(x, 'x', 'min');
      var xOrig = $x;
      var origAxes = util.parseAxisParam(axis, $x.shape);
      var axes = origAxes;
      var permutedAxes = axis_util.getAxesPermutation(axes, $x.rank);
      if (permutedAxes != null) {
          $x = $x.transpose(permutedAxes);
          axes = axis_util.getInnerMostAxes(axes.length, $x.rank);
      }
      var grad = function (dy, saved) {
          return gradForMinAndMax(dy, saved[1], saved[0], origAxes, permutedAxes);
      };
      var inputsToSave = [$x];
      var outputsToSave = [true];
      var res = engine_1.ENGINE.runKernelFunc(function (backend, save) {
          var y = backend.min($x, axes);
          save([xOrig, y]);
          return y;
      }, { x: $x }, grad, 'Min', { axes: axes }, inputsToSave, outputsToSave);
      if (keepDims) {
          var newShape = axis_util.expandShapeToKeepDim(res.shape, origAxes);
          res = res.reshape(newShape);
      }
      return res;
  }
  /**
   * Computes the maximum of elements across dimensions of a `tf.Tensor`.
   *
   * Reduces the input along the dimensions given in `axes`. Unless `keepDims`
   * is true, the rank of the `tf.Tensor` is reduced by 1 for each entry in
   * `axes`. If `keepDims` is true, the reduced dimensions are retained with
   * length 1. If `axes` has no entries, all dimensions are reduced, and an
   * `tf.Tensor` with a single element is returned.
   *
   * ```js
   * const x = tf.tensor1d([1, 2, 3]);
   *
   * x.max().print();  // or tf.max(x)
   * ```
   *
   * ```js
   * const x = tf.tensor2d([1, 2, 3, 4], [2, 2]);
   *
   * const axis = 1;
   * x.max(axis).print();  // or tf.max(x, axis)
   * ```
   *
   * @param x The input tensor.
   * @param axis The dimension(s) to reduce. By default it reduces
   *     all dimensions.
   * @param keepDims If true, retains reduced dimensions with size 1.
   */
  /** @doc {heading: 'Operations', subheading: 'Reduction'} */
  function max_(x, axis, keepDims) {
      if (axis === void 0) { axis = null; }
      if (keepDims === void 0) { keepDims = false; }
      var $x = tensor_util_env_1.convertToTensor(x, 'x', 'max');
      var xOrig = $x;
      var origAxes = util.parseAxisParam(axis, $x.shape);
      var axes = origAxes;
      var permutedAxes = axis_util.getAxesPermutation(axes, $x.rank);
      if (permutedAxes != null) {
          $x = $x.transpose(permutedAxes);
          axes = axis_util.getInnerMostAxes(axes.length, $x.rank);
      }
      var grad = function (dy, saved) {
          return gradForMinAndMax(dy, saved[1], saved[0], origAxes, permutedAxes);
      };
      var inputsToSave = [$x];
      var outputsToSave = [true];
      var res = engine_1.ENGINE.runKernelFunc(function (backend, save) {
          var y = backend.max($x, axes);
          save([xOrig, y]);
          return y;
      }, { x: $x }, grad, 'Max', { axes: axes }, inputsToSave, outputsToSave);
      if (keepDims) {
          var newShape = axis_util.expandShapeToKeepDim(res.shape, origAxes);
          res = res.reshape(newShape);
      }
      return res;
  }
  /**
   * Returns the indices of the minimum values along an `axis`.
   *
   * The result has the same shape as `input` with the dimension along `axis`
   * removed.
   *
   * ```js
   * const x = tf.tensor1d([1, 2, 3]);
   *
   * x.argMin().print();  // or tf.argMin(x)
   * ```
   *
   * ```js
   * const x = tf.tensor2d([1, 2, 4, 3], [2, 2]);
   *
   * const axis = 1;
   * x.argMin(axis).print();  // or tf.argMin(x, axis)
   * ```
   *
   * @param x The input tensor.
   * @param axis The dimension to reduce. Defaults to 0 (outer-most dimension).
   *
   */
  /** @doc {heading: 'Operations', subheading: 'Reduction'} */
  function argMin_(x, axis) {
      if (axis === void 0) { axis = 0; }
      var $x = tensor_util_env_1.convertToTensor(x, 'x', 'argMin');
      if (axis == null) {
          axis = 0;
      }
      var axes = util.parseAxisParam(axis, $x.shape);
      var permutedAxes = axis_util.getAxesPermutation(axes, $x.rank);
      if (permutedAxes != null) {
          $x = $x.transpose(permutedAxes);
          axes = axis_util.getInnerMostAxes(axes.length, $x.rank);
      }
      var grad = function (dy, saved) {
          var $x = saved[0];
          return { $x: function () { return tensor_ops_1.zerosLike($x); } };
      };
      return engine_1.ENGINE.runKernelFunc(function (backend, save) {
          var res = backend.argMin($x, axes[0]);
          save([$x]);
          return res;
      }, { $x: $x }, grad);
  }
  /**
   * Returns the indices of the maximum values along an `axis`.
   *
   * The result has the same shape as `input` with the dimension along `axis`
   * removed.
   *
   * ```js
   * const x = tf.tensor1d([1, 2, 3]);
   *
   * x.argMax().print();  // or tf.argMax(x)
   * ```
   *
   * ```js
   * const x = tf.tensor2d([1, 2, 4, 3], [2, 2]);
   *
   * const axis = 1;
   * x.argMax(axis).print();  // or tf.argMax(x, axis)
   * ```
   *
   * @param x The input tensor.
   * @param axis The dimension to reduce. Defaults to 0 (outer-most dimension).
   */
  /** @doc {heading: 'Operations', subheading: 'Reduction'} */
  function argMax_(x, axis) {
      if (axis === void 0) { axis = 0; }
      var $x = tensor_util_env_1.convertToTensor(x, 'x', 'argMax');
      if (axis == null) {
          axis = 0;
      }
      var axes = util.parseAxisParam(axis, $x.shape);
      var permutedAxes = axis_util.getAxesPermutation(axes, $x.rank);
      if (permutedAxes != null) {
          $x = $x.transpose(permutedAxes);
          axes = axis_util.getInnerMostAxes(axes.length, $x.rank);
      }
      var grad = function (dy, saved) {
          var $x = saved[0];
          return { x: function () { return tensor_ops_1.zerosLike($x); } };
      };
      var attrs = { axis: axes[0] };
      var inputsToSave = [$x];
      return engine_1.ENGINE.runKernelFunc(function (backend, save) {
          var res = backend.argMax($x, axes[0]);
          save([$x]);
          return res;
      }, { x: $x }, grad, 'ArgMax', attrs, inputsToSave);
  }
  /**
   * Computes the logical and of elements across dimensions of a `tf.Tensor`.
   *
   * Reduces the input along the dimensions given in `axes`. Unless `keepDims`
   * is true, the rank of the `tf.Tensor` is reduced by 1 for each entry in
   * `axes`. If `keepDims` is true, the reduced dimensions are retained with
   * length 1. If `axes` has no entries, all dimensions are reduced, and an
   * `tf.Tensor` with a single element is returned.
   *
   * ```js
   * const x = tf.tensor1d([1, 1, 1], 'bool');
   *
   * x.all().print();  // or tf.all(x)
   * ```
   *
   * ```js
   * const x = tf.tensor2d([1, 1, 0, 0], [2, 2], 'bool');
   *
   * const axis = 1;
   * x.all(axis).print();  // or tf.all(x, axis)
   * ```
   *
   * @param x The input tensor. Must be of dtype bool.
   * @param axis The dimension(s) to reduce. By default it reduces
   *     all dimensions.
   * @param keepDims If true, retains reduced dimensions with size 1.
   */
  /** @doc {heading: 'Operations', subheading: 'Reduction'} */
  function all_(x, axis, keepDims) {
      if (axis === void 0) { axis = null; }
      if (keepDims === void 0) { keepDims = false; }
      var $x = tensor_util_env_1.convertToTensor(x, 'x', 'all', 'bool');
      var origAxes = util.parseAxisParam(axis, $x.shape);
      var axes = origAxes;
      var permutedAxes = axis_util.getAxesPermutation(axes, $x.rank);
      if (permutedAxes != null) {
          $x = $x.transpose(permutedAxes);
          axes = axis_util.getInnerMostAxes(axes.length, $x.rank);
      }
      var res = engine_1.ENGINE.runKernelFunc(function (backend) { return backend.all($x, axes); }, { $x: $x });
      if (keepDims) {
          var newShape = axis_util.expandShapeToKeepDim(res.shape, origAxes);
          return res.reshape(newShape);
      }
      return res;
  }
  /**
   * Computes the logical or of elements across dimensions of a `tf.Tensor`.
   *
   * Reduces the input along the dimensions given in `axes`. Unless `keepDims`
   * is true, the rank of the `tf.Tensor` is reduced by 1 for each entry in
   * `axes`. If `keepDims` is true, the reduced dimensions are retained with
   * length 1. If `axes` has no entries, all dimensions are reduced, and an
   * `tf.Tensor` with a single element is returned.
   *
   * ```js
   * const x = tf.tensor1d([1, 1, 1], 'bool');
   *
   * x.any().print();  // or tf.any(x)
   * ```
   *
   * ```js
   * const x = tf.tensor2d([1, 1, 0, 0], [2, 2], 'bool');
   *
   * const axis = 1;
   * x.any(axis).print();  // or tf.any(x, axis)
   * ```
   *
   * @param x The input tensor. Must be of dtype bool.
   * @param axis The dimension(s) to reduce. By default it reduces
   *     all dimensions.
   * @param keepDims If true, retains reduced dimensions with size 1.
   */
  /** @doc {heading: 'Operations', subheading: 'Reduction'} */
  function any_(x, axis, keepDims) {
      if (axis === void 0) { axis = null; }
      if (keepDims === void 0) { keepDims = false; }
      var $x = tensor_util_env_1.convertToTensor(x, 'x', 'any', 'bool');
      var origAxes = util.parseAxisParam(axis, $x.shape);
      var axes = origAxes;
      var permutedAxes = axis_util.getAxesPermutation(axes, $x.rank);
      if (permutedAxes != null) {
          $x = $x.transpose(permutedAxes);
          axes = axis_util.getInnerMostAxes(axes.length, $x.rank);
      }
      var res = engine_1.ENGINE.runKernelFunc(function (backend) { return backend.any($x, axes); }, { $x: $x });
      if (keepDims) {
          var newShape = axis_util.expandShapeToKeepDim(res.shape, origAxes);
          return res.reshape(newShape);
      }
      return res;
  }
  /**
   * Calculates the mean and variance of `x`. The mean and variance are
   * calculated by aggregating the contents of `x` across `axes`. If `x` is
   * 1-D and `axes = [0]` this is just the mean and variance of a vector.
   *
   * @param x The input tensor.
   * @param axis The dimension(s) along with to compute mean and
   *     variance. By default it reduces all dimensions.
   * @param keepDims If true, the moments have the same dimensionality as the
   *     input.
   * @return An object with two keys: `mean` and `variance`.
   */
  /** @doc {heading: 'Operations', subheading: 'Normalization'} */
  function moments_(x, axis, keepDims) {
      if (axis === void 0) { axis = null; }
      if (keepDims === void 0) { keepDims = false; }
      x = tensor_util_env_1.convertToTensor(x, 'x', 'moments');
      var axes = util.parseAxisParam(axis, x.shape);
      var mean = x.mean(axes, keepDims);
      var keepDimsShape = mean.shape;
      if (!keepDims) {
          keepDimsShape = axis_util.expandShapeToKeepDim(mean.shape, axes);
      }
      var devSquared = x.toFloat().sub(mean.reshape(keepDimsShape)).square();
      var variance = devSquared.mean(axes, keepDims);
      return { mean: mean, variance: variance };
  }
  exports.all = operation_1.op({ all_: all_ });
  // tslint:disable-next-line:variable-name
  exports.any = operation_1.op({ any_: any_ });
  exports.argMax = operation_1.op({ argMax_: argMax_ });
  exports.argMin = operation_1.op({ argMin_: argMin_ });
  exports.logSumExp = operation_1.op({ logSumExp_: logSumExp_ });
  exports.max = operation_1.op({ max_: max_ });
  exports.mean = operation_1.op({ mean_: mean_ });
  exports.min = operation_1.op({ min_: min_ });
  exports.moments = operation_1.op({ moments_: moments_ });
  exports.sum = operation_1.op({ sum_: sum_ });
  exports.prod = operation_1.op({ prod_: prod_ });
  //# sourceMappingURL=reduction_ops.js.map