gx

chenyc

2025-02-12 ea42ff3ebee1eeb3fb29423aa848a249441db81c

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
/**
 * @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 { backend_util, env, ScatterNd } from '@tensorflow/tfjs-core';
import { ScatterProgram } from '../scatter_gpu';
import { ScatterPackedProgram } from '../scatter_packed_gpu';
import { reshape } from './Reshape';
export function scatterNd(args) {
    const { inputs, backend, attrs } = args;
    const { indices, updates } = inputs;
    const { shape } = attrs;
    const { sliceRank, numUpdates, sliceSize, strides, outputSize } = backend_util.calculateShapes(updates, indices, shape);
    const flattenShape = [outputSize / sliceSize, sliceSize];
    if (outputSize === 0) {
        return backend.makeTensorInfo(shape, indices.dtype);
    }
    const flattenIndices = reshape({ inputs: { x: indices }, backend, attrs: { shape: [numUpdates, sliceRank] } });
    const flattenX = reshape({ inputs: { x: updates }, backend, attrs: { shape: [numUpdates, sliceSize] } });
    const defaultValue = backend.makeTensorInfo([], 'float32', new Float32Array([0])); // scalar(0)
    let program;
    if (env().getBool('WEBGL_PACK')) {
        program = new ScatterPackedProgram(numUpdates, sliceRank, flattenIndices.shape.length, flattenX.shape.length, strides, flattenShape);
    }
    else {
        program = new ScatterProgram(numUpdates, sliceRank, flattenIndices.shape.length, flattenX.shape.length, strides, flattenShape);
    }
    const res = backend.runWebGLProgram(program, [flattenX, flattenIndices, defaultValue], flattenX.dtype);
    const reshaped = reshape({ inputs: { x: res }, backend, attrs: { shape } });
    backend.disposeIntermediateTensorInfo(flattenIndices);
    backend.disposeIntermediateTensorInfo(flattenX);
    backend.disposeIntermediateTensorInfo(res);
    backend.disposeIntermediateTensorInfo(defaultValue);
    return reshaped;
}
export const scatterNdConfig = {
    kernelName: ScatterNd,
    backendName: 'webgl',
    kernelFunc: scatterNd
};
//# sourceMappingURL=data:application/json;base64,eyJ2ZXJzaW9uIjozLCJmaWxlIjoiU2NhdHRlck5kLmpzIiwic291cmNlUm9vdCI6IiIsInNvdXJjZXMiOlsiLi4vLi4vLi4vLi4vLi4vLi4vdGZqcy1iYWNrZW5kLXdlYmdsL3NyYy9rZXJuZWxzL1NjYXR0ZXJOZC50cyJdLCJuYW1lcyI6W10sIm1hcHBpbmdzIjoiQUFBQTs7Ozs7Ozs7Ozs7Ozs7O0dBZUc7QUFFSCxPQUFPLEVBQUMsWUFBWSxFQUFFLEdBQUcsRUFBNEIsU0FBUyxFQUE4QyxNQUFNLHVCQUF1QixDQUFDO0FBRzFJLE9BQU8sRUFBQyxjQUFjLEVBQUMsTUFBTSxnQkFBZ0IsQ0FBQztBQUM5QyxPQUFPLEVBQUMsb0JBQW9CLEVBQUMsTUFBTSx1QkFBdUIsQ0FBQztBQUMzRCxPQUFPLEVBQUMsT0FBTyxFQUFDLE1BQU0sV0FBVyxDQUFDO0FBRWxDLE1BQU0sVUFBVSxTQUFTLENBQUMsSUFJekI7SUFDQyxNQUFNLEVBQUMsTUFBTSxFQUFFLE9BQU8sRUFBRSxLQUFLLEVBQUMsR0FBRyxJQUFJLENBQUM7SUFDdEMsTUFBTSxFQUFDLE9BQU8sRUFBRSxPQUFPLEVBQUMsR0FBRyxNQUFNLENBQUM7SUFDbEMsTUFBTSxFQUFDLEtBQUssRUFBQyxHQUFHLEtBQUssQ0FBQztJQUV0QixNQUFNLEVBQUMsU0FBUyxFQUFFLFVBQVUsRUFBRSxTQUFTLEVBQUUsT0FBTyxFQUFFLFVBQVUsRUFBQyxHQUN6RCxZQUFZLENBQUMsZUFBZSxDQUFDLE9BQU8sRUFBRSxPQUFPLEVBQUUsS0FBSyxDQUFDLENBQUM7SUFFMUQsTUFBTSxZQUFZLEdBQUcsQ0FBQyxVQUFVLEdBQUcsU0FBUyxFQUFFLFNBQVMsQ0FBQyxDQUFDO0lBRXpELElBQUksVUFBVSxLQUFLLENBQUMsRUFBRTtRQUNwQixPQUFPLE9BQU8sQ0FBQyxjQUFjLENBQUMsS0FBSyxFQUFFLE9BQU8sQ0FBQyxLQUFLLENBQUMsQ0FBQztLQUNyRDtJQUVELE1BQU0sY0FBYyxHQUFHLE9BQU8sQ0FDMUIsRUFBQyxNQUFNLEVBQUUsRUFBQyxDQUFDLEVBQUUsT0FBTyxFQUFDLEVBQUUsT0FBTyxFQUFFLEtBQUssRUFBRSxFQUFDLEtBQUssRUFBRSxDQUFDLFVBQVUsRUFBRSxTQUFTLENBQUMsRUFBQyxFQUFDLENBQUMsQ0FBQztJQUM5RSxNQUFNLFFBQVEsR0FBRyxPQUFPLENBQ3BCLEVBQUMsTUFBTSxFQUFFLEVBQUMsQ0FBQyxFQUFFLE9BQU8sRUFBQyxFQUFFLE9BQU8sRUFBRSxLQUFLLEVBQUUsRUFBQyxLQUFLLEVBQUUsQ0FBQyxVQUFVLEVBQUUsU0FBUyxDQUFDLEVBQUMsRUFBQyxDQUFDLENBQUM7SUFFOUUsTUFBTSxZQUFZLEdBQUcsT0FBTyxDQUFDLGNBQWMsQ0FDdkMsRUFBRSxFQUFFLFNBQVMsRUFBRSxJQUFJLFlBQVksQ0FBQyxDQUFDLENBQUMsQ0FBQyxDQUFDLENBQUMsQ0FBQyxDQUFFLFlBQVk7SUFDeEQsSUFBSSxPQUFPLENBQUM7SUFDWixJQUFJLEdBQUcsRUFBRSxDQUFDLE9BQU8sQ0FBQyxZQUFZLENBQUMsRUFBRTtRQUMvQixPQUFPLEdBQUcsSUFBSSxvQkFBb0IsQ0FDOUIsVUFBVSxFQUFFLFNBQVMsRUFBRSxjQUFjLENBQUMsS0FBSyxDQUFDLE1BQU0sRUFDbEQsUUFBUSxDQUFDLEtBQUssQ0FBQyxNQUFNLEVBQUUsT0FBTyxFQUFFLFlBQVksQ0FBQyxDQUFDO0tBQ25EO1NBQU07UUFDTCxPQUFPLEdBQUcsSUFBSSxjQUFjLENBQ3hCLFVBQVUsRUFBRSxTQUFTLEVBQUUsY0FBYyxDQUFDLEtBQUssQ0FBQyxNQUFNLEVBQ2xELFFBQVEsQ0FBQyxLQUFLLENBQUMsTUFBTSxFQUFFLE9BQU8sRUFBRSxZQUFZLENBQUMsQ0FBQztLQUNuRDtJQUNELE1BQU0sR0FBRyxHQUFHLE9BQU8sQ0FBQyxlQUFlLENBQy9CLE9BQU8sRUFBRSxDQUFDLFFBQVEsRUFBRSxjQUFjLEVBQUUsWUFBWSxDQUFDLEVBQUUsUUFBUSxDQUFDLEtBQUssQ0FBQyxDQUFDO0lBRXZFLE1BQU0sUUFBUSxHQUFHLE9BQU8sQ0FBQyxFQUFDLE1BQU0sRUFBRSxFQUFDLENBQUMsRUFBRSxHQUFHLEVBQUMsRUFBRSxPQUFPLEVBQUUsS0FBSyxFQUFFLEVBQUMsS0FBSyxFQUFDLEVBQUMsQ0FBQyxDQUFDO0lBRXRFLE9BQU8sQ0FBQyw2QkFBNkIsQ0FBQyxjQUFjLENBQUMsQ0FBQztJQUN0RCxPQUFPLENBQUMsNkJBQTZCLENBQUMsUUFBUSxDQUFDLENBQUM7SUFDaEQsT0FBTyxDQUFDLDZCQUE2QixDQUFDLEdBQUcsQ0FBQyxDQUFDO0lBQzNDLE9BQU8sQ0FBQyw2QkFBNkIsQ0FBQyxZQUFZLENBQUMsQ0FBQztJQUVwRCxPQUFPLFFBQVEsQ0FBQztBQUNsQixDQUFDO0FBRUQsTUFBTSxDQUFDLE1BQU0sZUFBZSxHQUFpQjtJQUMzQyxVQUFVLEVBQUUsU0FBUztJQUNyQixXQUFXLEVBQUUsT0FBTztJQUNwQixVQUFVLEVBQUUsU0FBa0M7Q0FDL0MsQ0FBQyIsInNvdXJjZXNDb250ZW50IjpbIi8qKlxuICogQGxpY2Vuc2VcbiAqIENvcHlyaWdodCAyMDIwIEdvb2dsZSBMTEMuIEFsbCBSaWdodHMgUmVzZXJ2ZWQuXG4gKiBMaWNlbnNlZCB1bmRlciB0aGUgQXBhY2hlIExpY2Vuc2UsIFZlcnNpb24gMi4wICh0aGUgXCJMaWNlbnNlXCIpO1xuICogeW91IG1heSBub3QgdXNlIHRoaXMgZmlsZSBleGNlcHQgaW4gY29tcGxpYW5jZSB3aXRoIHRoZSBMaWNlbnNlLlxuICogWW91IG1heSBvYnRhaW4gYSBjb3B5IG9mIHRoZSBMaWNlbnNlIGF0XG4gKlxuICogaHR0cDovL3d3dy5hcGFjaGUub3JnL2xpY2Vuc2VzL0xJQ0VOU0UtMi4wXG4gKlxuICogVW5sZXNzIHJlcXVpcmVkIGJ5IGFwcGxpY2FibGUgbGF3IG9yIGFncmVlZCB0byBpbiB3cml0aW5nLCBzb2Z0d2FyZVxuICogZGlzdHJpYnV0ZWQgdW5kZXIgdGhlIExpY2Vuc2UgaXMgZGlzdHJpYnV0ZWQgb24gYW4gXCJBUyBJU1wiIEJBU0lTLFxuICogV0lUSE9VVCBXQVJSQU5USUVTIE9SIENPTkRJVElPTlMgT0YgQU5ZIEtJTkQsIGVpdGhlciBleHByZXNzIG9yIGltcGxpZWQuXG4gKiBTZWUgdGhlIExpY2Vuc2UgZm9yIHRoZSBzcGVjaWZpYyBsYW5ndWFnZSBnb3Zlcm5pbmcgcGVybWlzc2lvbnMgYW5kXG4gKiBsaW1pdGF0aW9ucyB1bmRlciB0aGUgTGljZW5zZS5cbiAqID09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09XG4gKi9cblxuaW1wb3J0IHtiYWNrZW5kX3V0aWwsIGVudiwgS2VybmVsQ29uZmlnLCBLZXJuZWxGdW5jLCBTY2F0dGVyTmQsIFNjYXR0ZXJOZEF0dHJzLCBTY2F0dGVyTmRJbnB1dHMsIFRlbnNvckluZm99IGZyb20gJ0B0ZW5zb3JmbG93L3RmanMtY29yZSc7XG5cbmltcG9ydCB7TWF0aEJhY2tlbmRXZWJHTH0gZnJvbSAnLi4vYmFja2VuZF93ZWJnbCc7XG5pbXBvcnQge1NjYXR0ZXJQcm9ncmFtfSBmcm9tICcuLi9zY2F0dGVyX2dwdSc7XG5pbXBvcnQge1NjYXR0ZXJQYWNrZWRQcm9ncmFtfSBmcm9tICcuLi9zY2F0dGVyX3BhY2tlZF9ncHUnO1xuaW1wb3J0IHtyZXNoYXBlfSBmcm9tICcuL1Jlc2hhcGUnO1xuXG5leHBvcnQgZnVuY3Rpb24gc2NhdHRlck5kKGFyZ3M6IHtcbiAgaW5wdXRzOiBTY2F0dGVyTmRJbnB1dHMsXG4gIGJhY2tlbmQ6IE1hdGhCYWNrZW5kV2ViR0wsXG4gIGF0dHJzOiBTY2F0dGVyTmRBdHRyc1xufSk6IFRlbnNvckluZm8ge1xuICBjb25zdCB7aW5wdXRzLCBiYWNrZW5kLCBhdHRyc30gPSBhcmdzO1xuICBjb25zdCB7aW5kaWNlcywgdXBkYXRlc30gPSBpbnB1dHM7XG4gIGNvbnN0IHtzaGFwZX0gPSBhdHRycztcblxuICBjb25zdCB7c2xpY2VSYW5rLCBudW1VcGRhdGVzLCBzbGljZVNpemUsIHN0cmlkZXMsIG91dHB1dFNpemV9ID1cbiAgICAgIGJhY2tlbmRfdXRpbC5jYWxjdWxhdGVTaGFwZXModXBkYXRlcywgaW5kaWNlcywgc2hhcGUpO1xuXG4gIGNvbnN0IGZsYXR0ZW5TaGFwZSA9IFtvdXRwdXRTaXplIC8gc2xpY2VTaXplLCBzbGljZVNpemVdO1xuXG4gIGlmIChvdXRwdXRTaXplID09PSAwKSB7XG4gICAgcmV0dXJuIGJhY2tlbmQubWFrZVRlbnNvckluZm8oc2hhcGUsIGluZGljZXMuZHR5cGUpO1xuICB9XG5cbiAgY29uc3QgZmxhdHRlbkluZGljZXMgPSByZXNoYXBlKFxuICAgICAge2lucHV0czoge3g6IGluZGljZXN9LCBiYWNrZW5kLCBhdHRyczoge3NoYXBlOiBbbnVtVXBkYXRlcywgc2xpY2VSYW5rXX19KTtcbiAgY29uc3QgZmxhdHRlblggPSByZXNoYXBlKFxuICAgICAge2lucHV0czoge3g6IHVwZGF0ZXN9LCBiYWNrZW5kLCBhdHRyczoge3NoYXBlOiBbbnVtVXBkYXRlcywgc2xpY2VTaXplXX19KTtcblxuICBjb25zdCBkZWZhdWx0VmFsdWUgPSBiYWNrZW5kLm1ha2VUZW5zb3JJbmZvKFxuICAgICAgW10sICdmbG9hdDMyJywgbmV3IEZsb2F0MzJBcnJheShbMF0pKTsgIC8vIHNjYWxhcigwKVxuICBsZXQgcHJvZ3JhbTtcbiAgaWYgKGVudigpLmdldEJvb2woJ1dFQkdMX1BBQ0snKSkge1xuICAgIHByb2dyYW0gPSBuZXcgU2NhdHRlclBhY2tlZFByb2dyYW0oXG4gICAgICAgIG51bVVwZGF0ZXMsIHNsaWNlUmFuaywgZmxhdHRlbkluZGljZXMuc2hhcGUubGVuZ3RoLFxuICAgICAgICBmbGF0dGVuWC5zaGFwZS5sZW5ndGgsIHN0cmlkZXMsIGZsYXR0ZW5TaGFwZSk7XG4gIH0gZWxzZSB7XG4gICAgcHJvZ3JhbSA9IG5ldyBTY2F0dGVyUHJvZ3JhbShcbiAgICAgICAgbnVtVXBkYXRlcywgc2xpY2VSYW5rLCBmbGF0dGVuSW5kaWNlcy5zaGFwZS5sZW5ndGgsXG4gICAgICAgIGZsYXR0ZW5YLnNoYXBlLmxlbmd0aCwgc3RyaWRlcywgZmxhdHRlblNoYXBlKTtcbiAgfVxuICBjb25zdCByZXMgPSBiYWNrZW5kLnJ1bldlYkdMUHJvZ3JhbShcbiAgICAgIHByb2dyYW0sIFtmbGF0dGVuWCwgZmxhdHRlbkluZGljZXMsIGRlZmF1bHRWYWx1ZV0sIGZsYXR0ZW5YLmR0eXBlKTtcblxuICBjb25zdCByZXNoYXBlZCA9IHJlc2hhcGUoe2lucHV0czoge3g6IHJlc30sIGJhY2tlbmQsIGF0dHJzOiB7c2hhcGV9fSk7XG5cbiAgYmFja2VuZC5kaXNwb3NlSW50ZXJtZWRpYXRlVGVuc29ySW5mbyhmbGF0dGVuSW5kaWNlcyk7XG4gIGJhY2tlbmQuZGlzcG9zZUludGVybWVkaWF0ZVRlbnNvckluZm8oZmxhdHRlblgpO1xuICBiYWNrZW5kLmRpc3Bvc2VJbnRlcm1lZGlhdGVUZW5zb3JJbmZvKHJlcyk7XG4gIGJhY2tlbmQuZGlzcG9zZUludGVybWVkaWF0ZVRlbnNvckluZm8oZGVmYXVsdFZhbHVlKTtcblxuICByZXR1cm4gcmVzaGFwZWQ7XG59XG5cbmV4cG9ydCBjb25zdCBzY2F0dGVyTmRDb25maWc6IEtlcm5lbENvbmZpZyA9IHtcbiAga2VybmVsTmFtZTogU2NhdHRlck5kLFxuICBiYWNrZW5kTmFtZTogJ3dlYmdsJyxcbiAga2VybmVsRnVuYzogc2NhdHRlck5kIGFzIHVua25vd24gYXMgS2VybmVsRnVuY1xufTtcbiJdfQ==