Function sparseReshape

tfjs-backend-cpu/src/kernels/SparseReshape.ts:24–58

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

import {sparseReshapeImpl} from './SparseReshape_impl';

export function sparseReshape(
    args: {inputs: SparseReshapeInputs, backend: MathBackendCPU}):
    [TensorInfo, TensorInfo] {
  const {inputs, backend} = args;
  const {inputIndices, inputShape, newShape} = inputs;
  if (inputIndices.shape.length !== 2) {
    throw new Error(`Input indices should be a matrix but received shape
        ${inputIndices.shape}`);
  }
  if (inputShape.shape.length !== 1) {
    throw new Error(`Input shape should be a vector but received shape
        ${inputShape.shape}`);
  }

  if (newShape.shape.length !== 1) {
    throw new Error(
        `Target shape should be a vector but received shape ${newShape.shape}`);
  }

  const $inputShape =
      Array.from(backend.data.get(inputShape.dataId).values as TypedArray);
  const $inputIndices =
      backend.data.get(inputIndices.dataId).values as TypedArray;
  const targetShape =
      Array.from(backend.data.get(newShape.dataId).values as TypedArray);

  const [newIndices, indicesShape, outputShape] = sparseReshapeImpl(
      $inputIndices, inputIndices.shape, inputIndices.dtype, $inputShape,
      targetShape);
  return [
    backend.makeTensorInfo(indicesShape, inputIndices.dtype, newIndices),
    backend.makeTensorInfo(
        [outputShape.length], newShape.dtype, new Int32Array(outputShape)),
  ];
}

export const sparseReshapeConfig: KernelConfig = {
  kernelName: SparseReshape,