Function sparseSegmentSum

tfjs-backend-webgpu/src/kernels/SparseSegmentSum.ts:23–30

(
    args: {inputs: SparseSegmentSumInputs, backend: WebGPUBackend})

import {sparseSegmentReduce} from '../kernel_utils/sparse_segment_reduce';

export function sparseSegmentSum(
    args: {inputs: SparseSegmentSumInputs, backend: WebGPUBackend}):
    TensorInfo {
  const {inputs, backend} = args;
  const {data, indices, segmentIds} = inputs;

  return sparseSegmentReduce(data, indices, segmentIds, true, backend);
}

export const sparseSegmentSumConfig: KernelConfig = {
  kernelName: SparseSegmentSum,