Function sparseFillEmptyRows

tfjs-backend-cpu/src/kernels/SparseFillEmptyRows.ts:24–70

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

import {sparseFillEmptyRowsImpl} from './SparseFillEmptyRows_impl';

export function sparseFillEmptyRows(args: {
  inputs: SparseFillEmptyRowsInputs,
  backend: MathBackendCPU
}): [TensorInfo, TensorInfo, TensorInfo, TensorInfo] {
  const {inputs, backend} = args;
  const {indices, values, denseShape, defaultValue} = inputs;
  if (denseShape.shape.length !== 1) {
    throw new Error(`Dense shape must be a vector, saw:
        ${denseShape.shape}`);
  }
  if (indices.shape.length !== 2) {
    throw new Error(`Indices must be a matrix, saw:
        ${indices.shape}`);
  }
  if (values.shape.length !== 1) {
    throw new Error(`Values must be a vector, saw:
        ${values.shape}`);
  }
  if (defaultValue.shape.length !== 0) {
    throw new Error(`Default value must be a scalar, saw:
        ${defaultValue.shape}`);
  }

  const $indices = backend.data.get(indices.dataId).values as TypedArray;
  const $values = backend.data.get(values.dataId).values as TypedArray;
  const $denseShape = backend.data.get(denseShape.dataId).values as TypedArray;
  const $defaultValue =
      backend.data.get(defaultValue.dataId).values[0] as number;

  const [outputIndices, outputIndicesShape, outputValues,
         emptyRowIndicator, reverseIndexMap] =
      sparseFillEmptyRowsImpl(
          $indices, indices.shape, indices.dtype, $values, values.dtype,
          $denseShape, $defaultValue);
  return [
    backend.makeTensorInfo(outputIndicesShape, indices.dtype, outputIndices),
    backend.makeTensorInfo(
        [outputIndicesShape[0]], values.dtype, outputValues),
    backend.makeTensorInfo(
        [emptyRowIndicator.length], 'bool',
        new Uint8Array(
            emptyRowIndicator.map((value: boolean) => Number(value)))),
    backend.makeTensorInfo(
        [reverseIndexMap.length], indices.dtype,
        new Int32Array(reverseIndexMap)),
  ];
}

export const sparseFillEmptyRowsConfig: KernelConfig = {
  kernelName: SparseFillEmptyRows,