license import { BackendValues , DataType , TensorBuffer , TypedArray , util } from '@tensorflow/tfjs-core' ; export function uniqueImpl ( values , axis , shape , dtype ) { const $axis = util . parseAxisParam ( axis , shape ) [ 0 ] ; const newShape = [ 1 , shape [ 0 ] , 1 ] ; for ( let i = 0 ; i < $axis ; i ++ ) { newShape [ 0 ] *= shape [ i ] ; } newShape [ 1 ] = shape [ $axis ] ; for ( let i = $axis + 1 ; i < shape . length ; i ++ ) { newShape [ 2 ] *= shape [ i ] ; } const uniqueElements = { } ; const indices = new Int32Array ( shape [ $axis ] ) ; const inputBuffer = new TensorBuffer ( newShape , dtype , values as TypedArray ) ; const uniqueIndices = [ ] ; const is1DTensor = newShape [ 0 ] === 1 && newShape [ 2 ] === 1 ; for ( let i = 0 ; i < shape [ $axis ] ; i ++ ) { let element ; if ( is1DTensor ) { element = values [ i ] . toString ( ) ; } else { const axisValues = [ ] ; for ( let m = 0 ; m < newShape [ 0 ] ; m ++ ) { for ( let n = 0 ; n < newShape [ 2 ] ; n ++ ) { axisValues . push ( inputBuffer . get ( m , i , n ) ) ; } } element = axisValues . join ( ',' ) ; } if ( uniqueElements [ element ] !== undefined ) { indices [ i ] = uniqueElements [ element ] ; } else { const uniqueIndex = Object . keys ( uniqueElements ) . length ; uniqueElements [ element ] = uniqueIndex ; indices [ i ] = uniqueIndex ; uniqueIndices . push ( i ) ; } } const outputTmpShape = newShape . slice ( ) ; outputTmpShape [ 1 ] = Object . keys ( uniqueElements ) . length ; const outputBuffer = new TensorBuffer ( outputTmpShape , dtype ) ; uniqueIndices . forEach ( ( uniqueElementIndex , i ) => { for ( let m = 0 ; m < newShape [ 0 ] ; m ++ ) { for ( let n = 0 ; n < newShape [ 2 ] ; n ++ ) { outputBuffer . set ( inputBuffer . get ( m , uniqueElementIndex , n ) , m , i , n ) ; } } } ) ; const outputShape = shape . slice ( ) ; outputShape [ $axis ] = outputTmpShape [ 1 ] ; return { outputValues : outputBuffer . values as BackendValues , outputShape , indices , } ; }