Function stringSplitImpl

tfjs-backend-cpu/src/kernels/StringSplit_impl.ts:64–100

(
    input: Uint8Array[], delimiter: Uint8Array,
    skipEmpty: boolean)

}

export function stringSplitImpl(
    input: Uint8Array[], delimiter: Uint8Array,
    skipEmpty: boolean): [TypedArray, Uint8Array[], [number, number]] {
  const batchSize = input.length;

  // Empty delimiter means split the input character by character.
  const tokens: Uint8Array[] = [];

  let outputSize = 0;
  let maxNumEntries = 0;
  const numIndices: number[] = new Array(batchSize);
  for (let i = 0; i < batchSize; ++i) {
    const prevTokensLength = tokens.length;
    split(input[i], delimiter, skipEmpty, tokens);
    const nEntries = tokens.length - prevTokensLength;
    numIndices[i] = nEntries;
    outputSize += nEntries;
    maxNumEntries = Math.max(maxNumEntries, nEntries);
  }

  const indices = util.getArrayFromDType('int32', outputSize * 2) as TypedArray;
  const values: Uint8Array[] = new Array(outputSize);
  const shape: [number, number] = [batchSize, maxNumEntries];

  let c = 0;
  for (let i = 0; i < batchSize; ++i) {
    for (let j = 0; j < numIndices[i]; ++j) {
      // indices is a 2d tensor with shape of [outputSize, 2]
      indices[c * 2] = i;
      indices[c * 2 + 1] = j;
      values[c] = tokens[c];
      ++c;
    }
  }

  return [indices, values, shape];
}