INDArray bucketData ( INDArray query ) { INDArray mask = bucket ( query ) ; int nRes = mask . sum ( 0 ) . getInt ( 0 ) ; INDArray res = Nd4j . create ( new int [ ] { nRes , inDimension } ) ; int j = 0 ; for ( int i = 0 ; i < nRes ; i ++ ) { while ( mask . getInt ( j ) == 0 && j < mask . length ( ) - 1 ) { j += 1 ; } if ( mask . getInt ( j ) == 1 ) res . putRow ( i , indexData . getRow ( j ) ) ; j += 1 ; } return res ; }
data elements in the same entropy bucket as the query