Function sparseFillEmptyRows

tfjs-backend-wasm/src/kernels/SparseFillEmptyRows.ts:50–147

(args: {
  backend: BackendWasm,
  inputs: SparseFillEmptyRowsInputs,
})

}

export function sparseFillEmptyRows(args: {
  backend: BackendWasm,
  inputs: SparseFillEmptyRowsInputs,
}): [TensorInfo, TensorInfo, TensorInfo, TensorInfo] {
  const {backend, inputs} = args;
  const {indices, values, denseShape, defaultValue} = inputs;

  const indicesCount = indices.shape[0];
  const rank = indices.shape[1];
  const denseRows = backend.readSync(denseShape.dataId)[0] as number;

  // Set output size to maximum possible and resize later (actual result
  // might be smaller).
  const maxOutputIndicesShape = [indicesCount + denseRows, rank];

  const indicesId = backend.dataIdMap.get(indices.dataId).id;
  const valuesId = backend.dataIdMap.get(values.dataId).id;
  const defaultValueId = backend.dataIdMap.get(defaultValue.dataId).id;

  const outputIndices =
      backend.makeOutput(maxOutputIndicesShape, indices.dtype);
  const outputIndicesId = backend.dataIdMap.get(outputIndices.dataId).id;

  const outputValues =
      backend.makeOutput(maxOutputIndicesShape.slice(0, 1), values.dtype);
  const outputValuesId = backend.dataIdMap.get(outputValues.dataId).id;

  const emptyRowIndicator = backend.makeOutput([denseRows], 'bool');
  const emptyRowIndicatorId =
      backend.dataIdMap.get(emptyRowIndicator.dataId).id;

  const reverseIndexMap = backend.makeOutput([indicesCount], indices.dtype);
  const reverseIndexMapId = backend.dataIdMap.get(reverseIndexMap.dataId).id;

  const exceptionValues = backend.makeOutput([4], 'int32');
  const exceptionValuesId = backend.dataIdMap.get(exceptionValues.dataId).id;

  const outputRows = wasmSparseFillEmptyRows(
      indicesId, valuesId, CppDType[values.dtype], indicesCount, denseRows,
      rank, defaultValueId, outputIndicesId, outputValuesId,
      emptyRowIndicatorId, reverseIndexMapId, exceptionValuesId);

  const exceptionValuesArray =
      backend.readSync(exceptionValues.dataId) as Int32Array;

  let exceptionMessage: string;
  switch (exceptionValuesArray[0]) {
    case 1: {
      exceptionMessage =
          backend_util.getSparseFillEmptyRowsIndicesDenseShapeMismatch(
              exceptionValuesArray[1]);
      break;
    }
    case 2: {
      exceptionMessage =
          backend_util.getSparseFillEmptyRowsNegativeIndexErrorMessage(
              exceptionValuesArray[1], exceptionValuesArray[2]);
      break;