@tensorflow/tfjs-core
Version:
Hardware-accelerated JavaScript library for machine intelligence
53 lines • 2.29 kB
JavaScript
/**
* @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 { UnsortedSegmentSum } from '../kernel_names';
import { expandDims } from '../ops/expand_dims';
import { gather } from '../ops/gather';
import { greaterEqual } from '../ops/greater_equal';
import { logicalAnd } from '../ops/logical_and';
import { maximum } from '../ops/maximum';
import { ones } from '../ops/ones';
import { scalar } from '../ops/scalar';
import { where } from '../ops/where';
import { zerosLike } from '../ops/zeros_like';
export const unsortedSegmentSumGradConfig = {
kernelName: UnsortedSegmentSum,
inputsToSave: ['segmentIds'],
gradFunc: (dy, saved) => {
const [segmentIds] = saved;
const derX = () => {
return gatherDropNegatives(dy, segmentIds);
};
return { x: derX };
}
};
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
const zeroClippedIndices = maximum(indices, zerosLike(indices));
const gathered = gather(x, zeroClippedIndices);
let isPositive = greaterEqual(indices, scalar(0, 'int32'));
const numIters = gathered.rank - isPositive.rank;
for (let i = 0; i < numIters; ++i) {
isPositive = expandDims(isPositive, i + 1);
}
isPositive = logicalAnd(isPositive, ones(gathered.shape, 'bool'));
const zeroSlice = zerosLike(gathered);
return where(isPositive, gathered, zeroSlice);
}
//# sourceMappingURL=UnsortedSegmentSum_grad.js.map