Function sparseSegmentMean

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

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

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

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

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

export const sparseSegmentMeanConfig: KernelConfig = {
  kernelName: SparseSegmentMean,