"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 tensor_util_env_1 = require("../tensor_util_env");
var util_1 = require("../util");
var array_ops_1 = require("./array_ops");
var axis_util_1 = require("./axis_util");
var binary_ops_1 = require("./binary_ops");
var compare_1 = require("./compare");
var logical_ops_1 = require("./logical_ops");
var operation_1 = require("./operation");
var segment_util_1 = require("./segment_util");
var tensor_ops_1 = require("./tensor_ops");
/**
 * Computes the sum along segments of a `tf.Tensor`.
 *
 * ```js
 * const x = tf.tensor1d([1, 2, 3, 4]);
 * const segmentIds = tf.tensor1d([1, 2, 0, 1], 'int32');
 * const numSegments = 3;
 *
 * x.unsortedSegmentSum(segmentIds, numSegments).print()
 * //or tf.unsortedSegmentSum(x, segmentIds, numSegments)
 * ```
 * @param x The `tf.Tensor` that will be summed along its segments.
 * @param segmentIds A `tf.Tensor1D` whose rank is equal to the rank of `x`'s
 * dimension along the `axis`.  Maps each element of `x` to a segment.
 * @param numSegments The number of distinct `segmentIds`.
 */
/** @doc {heading: 'Operations', subheading: 'Segment'} */
function unsortedSegmentSum_(x, segmentIds, numSegments) {
    var $x = tensor_util_env_1.convertToTensor(x, 'x', 'unsortedSegmentSum');
    var $segmentIds = tensor_util_env_1.convertToTensor(segmentIds, 'segmentIds', 'unsortedSegmentSum', 'int32');
    util_1.assert(util_1.isInt(numSegments), function () { return 'numSegments must be of dtype int'; });
    var gradFunc = function (dy, saved) {
        var $segmentIds = saved[0];
        var derX = function () {
            return gatherDropNegatives(dy, $segmentIds);
        };
        return { $x: derX };
    };
    return engine_1.ENGINE.runKernelFunc(function (backend, save) {
        var res = backend.unsortedSegmentSum($x, $segmentIds, numSegments);
        save([$segmentIds]);
        return res;
    }, { $x: $x }, gradFunc);
}
/**
 * Gather slices from tensor `x`'s axis `axis` according to `indices`.
 *
 * ```js
 * const x = tf.tensor1d([1, 2, 3, 4]);
 * const indices = tf.tensor1d([1, 3, 3], 'int32');
 *
 * x.gather(indices).print();
 * ```
 *
 * ```js
 * const x = tf.tensor2d([1, 2, 3, 4], [2, 2]);
 * const indices = tf.tensor1d([1, 1, 0], 'int32');
 *
 * x.gather(indices).print();
 * ```
 * @param x The input tensor whose slices to be gathered.
 * @param indices The indices of the values to extract.
 * @param axis The axis over which to select values. Defaults to 0.
 */
/** @doc {heading: 'Tensors', subheading: 'Slicing and Joining'} */
function gather_(x, indices, axis) {
    if (axis === void 0) { axis = 0; }
    var $x = tensor_util_env_1.convertToTensor(x, 'x', 'gather');
    var $indices = tensor_util_env_1.convertToTensor(indices, 'indices', 'gather', 'int32');
    axis = util_1.parseAxisParam(axis, $x.shape)[0];
    var shapeInfo = segment_util_1.collectGatherOpShapeInfo($x, $indices, axis);
    var grad = function (dy, saved) {
        var $indices = saved[0];
        var derX = function () {
            var paramsShape = $x.shape;
            var indicesSize = $indices.size;
            var outerShape = paramsShape.slice(0, axis);
            var outerDims = outerShape.length;
            var innerShape = paramsShape.slice(axis, paramsShape.length).slice(1);
            var innerDims = innerShape.length;
            var outerAxesIndices = arrayRange(0, outerDims);
            var innerAxesIndices = arrayRange(outerDims + 1, outerDims + 1 + innerDims);
            var valuesShape = arrayConcat([outerShape, [indicesSize], innerShape]);
            var values = dy.reshape(valuesShape);
            var reshapedIndices = $indices.reshape([indicesSize]);
            var transposeDims = arrayConcat([[outerDims], outerAxesIndices, innerAxesIndices]);
            var valuesTranspose = values.transpose(transposeDims);
            var paramsGrad = exports.unsortedSegmentSum(valuesTranspose, reshapedIndices, $x.shape[axis]);
            var invertTransposeDims = axis_util_1.getUndoAxesPermutation(transposeDims);
            paramsGrad = paramsGrad.transpose(invertTransposeDims);
            return paramsGrad;
        };
        return { x: derX, indices: function () { return $indices; } };
    };
    return (engine_1.ENGINE.runKernelFunc(function (backend, save) {
        var res = backend.gather($x, $indices.flatten(), axis);
        save([$indices]);
        return res;
    }, { x: $x, indices: $indices }, grad, 'Gather', { axis: axis }))
        .reshape(shapeInfo.outputShape);
}
function arrayRange(start, stop) {
    var result = [];
    for (var i = start; i < stop; ++i) {
        result.push(i);
    }
    return result;
}
function arrayConcat(arrays) {
    var result = [];
    for (var i = 0; i < arrays.length; ++i) {
        for (var j = 0; j < arrays[i].length; ++j) {
            result.push(arrays[i][j]);
        }
    }
    return result;
}
function gatherDropNegatives(x, indices) {
    // Helper function for unsorted segment ops. Gathers params for
    // positive segment ids and gathers 0 for inputs with negative segment id.
    // Mirrors _GatherDropNegatives from tensorflow/python/ops/math_grad.py
    var zeroClippedIndices = binary_ops_1.maximum(indices, tensor_ops_1.zerosLike(indices));
    var gathered = exports.gather(x, zeroClippedIndices);
    var isPositive = compare_1.greaterEqual(indices, tensor_ops_1.scalar(0, 'int32'));
    var numIters = gathered.rank - isPositive.rank;
    for (var i = 0; i < numIters; ++i) {
        isPositive = array_ops_1.expandDims(isPositive, i + 1);
    }
    isPositive = logical_ops_1.logicalAnd(isPositive, tensor_ops_1.ones(gathered.shape, 'bool'));
    var zeroSlice = tensor_ops_1.zerosLike(gathered);
    return logical_ops_1.where(isPositive, gathered, zeroSlice);
}
exports.gather = operation_1.op({ gather_: gather_ });
exports.unsortedSegmentSum = operation_1.op({ unsortedSegmentSum_: unsortedSegmentSum_ });
//# sourceMappingURL=segment_ops.js.map