Function tensorScatterUpdate

tfjs-backend-wasm/src/kernels/TensorScatterUpdate.ts:45–79

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

}

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

  const out = backend.makeOutput(tensor.shape, tensor.dtype);
  if (util.sizeFromShape(tensor.shape) === 0) {
    return out;
  }

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

  const indicesData = backend.dataIdMap.get(indices.dataId);
  const indicesId = indicesData.id;

  const updatesData = backend.dataIdMap.get(updates.dataId);
  const updatesId = updatesData.id;

  const tensorData = backend.dataIdMap.get(tensor.dataId);
  const tensorId = tensorData.id;

  const stridesBytes = new Uint8Array(new Int32Array(strides).buffer);

  const outId = backend.dataIdMap.get(out.dataId).id;
  wasmTensorScatterUpdate(
      indicesId, updatesId, CppDType[updates.dtype], sliceRank, numUpdates,
      sliceSize, stridesBytes, outputSize, outId, tensorId);

  return out;
}

export const tensorScatterUpdateConfig: KernelConfig = {
  kernelName: TensorScatterUpdate,