@tensorflow/tfjs-core
Version:
Hardware-accelerated JavaScript library for machine intelligence
154 lines • 6.76 kB
JavaScript
"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