Function tensorScatterUpdate

tfjs-backend-cpu/src/kernels/TensorScatterUpdate.ts:24–43

(args: {
  inputs: TensorScatterUpdateInputs,
  backend: MathBackendCPU,
  attrs: TensorScatterUpdateAttrs
})

import {scatterImpl} from './Scatter_impl';

export function tensorScatterUpdate(args: {
  inputs: TensorScatterUpdateInputs,
  backend: MathBackendCPU,
  attrs: TensorScatterUpdateAttrs
}): TensorInfo {
  const {inputs, backend} = args;
  const {tensor, indices, updates} = inputs;

  const {sliceRank, numUpdates, sliceSize, strides, outputSize} =
      backend_util.calculateShapes(updates, indices, tensor.shape);
  const sumDupeIndices = false;

  const indicesBuf = backend.bufferSync<Rank, 'int32'>(indices);
  const updatesBuf = backend.bufferSync<Rank, 'int32'|'float32'>(updates);
  const tensorBuf = backend.bufferSync<Rank, 'int32'|'float32'>(tensor);
  const outBuf = scatterImpl(
      indicesBuf, updatesBuf, tensor.shape, outputSize, sliceSize, numUpdates,
      sliceRank, strides, tensorBuf, sumDupeIndices);
  return backend.makeTensorInfo(tensor.shape, outBuf.dtype, outBuf.values);
}

export const tensorScatterUpdateConfig: KernelConfig = {
  kernelName: TensorScatterUpdate,