license import { KernelConfig , KernelFunc , StringSplit , StringSplitAttrs , StringSplitInputs , TensorInfo } from '@tensorflow/tfjs-core' ; import { MathBackendCPU } from '../backend_cpu' ; import { stringSplitImpl } from './StringSplit_impl' ; export function stringSplit ( args ) { const { inputs , backend , attrs } = args ; const { skipEmpty } = attrs ; const { input , delimiter } = inputs ; if ( input . dtype !== 'string' ) { throw new Error ( 'Input must be of datatype string' ) ; } if ( input . shape . length !== 1 ) { throw new Error ( `Input must be a vector, got shape: ${ input . shape }` ) ; } if ( delimiter . shape . length !== 0 ) { throw new Error ( `Delimiter must be a scalar, got shape: ${ delimiter . shape }` ) ; } const $input = backend . data . get ( input . dataId ) . values as Uint8Array [ ] ; const $delimiter = backend . data . get ( delimiter . dataId ) . values [ 0 ] as Uint8Array ; const [ indices , values , shape ] = stringSplitImpl ( $input , $delimiter , skipEmpty ) ; const outputSize = values . length ; return [ backend . makeTensorInfo ( [ outputSize , 2 ] , 'int32' , indices ) , backend . makeTensorInfo ( [ outputSize ] , 'string' , values ) , backend . makeTensorInfo ( [ 2 ] , 'int32' , new Int32Array ( shape ) ) ] ; } export const stringSplitConfig = { kernelName : StringSplit , backendName : 'cpu' , kernelFunc : stringSplit as { } as KernelFunc , } ;