Function sparseToDense

tfjs-backend-webgl/src/kernels/SparseToDense.ts:26–61

(args: {
  inputs: SparseToDenseInputs,
  backend: MathBackendWebGL,
  attrs: SparseToDenseAttrs
})

import {reshape} from './Reshape';

export function sparseToDense(args: {
  inputs: SparseToDenseInputs,
  backend: MathBackendWebGL,
  attrs: SparseToDenseAttrs
}): TensorInfo {
  const {inputs, backend, attrs} = args;
  const {sparseIndices, sparseValues, defaultValue} = inputs;
  const {outputShape} = attrs;

  const {sliceRank, numUpdates, sliceSize, strides, outputSize} =
      backend_util.calculateShapes(sparseValues, sparseIndices, outputShape);
  const sumDupeIndices = false;

  if (sparseValues.dtype === 'string') {
    const indicesBuf = backend.bufferSync<Rank, 'int32'>(sparseIndices);
    const updatesBuf = backend.bufferSync<Rank, 'string'>(sparseValues);
    const $defaultValue = util.decodeString(
        backend.readSync(defaultValue.dataId)[0] as Uint8Array);
    const outBuf = scatterImplCPU(
        indicesBuf, updatesBuf, outputShape, outputSize, sliceSize, numUpdates,
        sliceRank, strides, $defaultValue, sumDupeIndices);
    return backend.makeTensorInfo(outputShape, outBuf.dtype, outBuf.values);
  }
  const program = new ScatterProgram(
      numUpdates, sliceRank, sparseIndices.shape.length,
      sparseValues.shape.length, strides, [outputSize, 1], sumDupeIndices);

  const res = backend.runWebGLProgram(
      program, [sparseValues, sparseIndices, defaultValue], sparseValues.dtype);

  const reshaped =
      reshape({inputs: {x: res}, backend, attrs: {shape: outputShape}});

  backend.disposeIntermediateTensorInfo(res);
  return reshaped;
}

export const sparseToDenseConfig: KernelConfig = {
  kernelName: SparseToDense,