license import * as tf from '@tensorflow/tfjs-core' ; import { describeWithFlags } from '@tensorflow/tfjs-core/dist/jasmine_util' ; import { WEBGL_ENVS } from '../backend_webgl_test_registry' ; describeWithFlags ( 'TopK' , WEBGL_ENVS , ( ) => { it ( 'handles packed inputs' , async ( ) => { const a = tf . tensor1d ( [ 1 , 1 , 2 , 1 , 2 , 1 , 1 , 1 , 2 , 2 , 1 , 1 , 2 , 1 , 1 , 2 , 1 , 1 , 2 , 1 , 1 , 2 , 1 , 2 , 1 , 1 , 1 , 2 , 1 , 1 , 2 , 1 , 1 , 2 , 1 , 1 , 1 , 1 , 2 , 1 , 2 , 2 , 2 , 2 , 1 , 2 , 1 , 2 , 2 , 2 , 1 , 2 , 2 , 1 , 2 , 2 , 1 , 2 , 1 , 1 , 1 , 2 , 1 , 2 , 1 ] ) ; tf . env ( ) . set ( 'WEBGL_PACK' , true ) ; const aPacked = tf . addN ( [ a , tf . zeros ( a . shape ) ] ) ; const k = a . shape [ 0 ] ; const { values , indices } = tf . topk ( aPacked , k ) ; expect ( values . shape ) . toEqual ( [ k ] ) ; expect ( indices . shape ) . toEqual ( [ k ] ) ; expect ( values . dtype ) . toBe ( 'float32' ) ; expect ( indices . dtype ) . toBe ( 'int32' ) ; tf . test_util . expectArraysEqual ( await values . data ( ) , [ 2 , 2 , 2 , 2 , 2 , 2 , 2 , 2 , 2 , 2 , 2 , 2 , 2 , 2 , 2 , 2 , 2 , 2 , 2 , 2 , 2 , 2 , 2 , 2 , 2 , 2 , 2 , 2 , 1 , 1 , 1 , 1 , 1 , 1 , 1 , 1 , 1 , 1 , 1 , 1 , 1 , 1 , 1 , 1 , 1 , 1 , 1 , 1 , 1 , 1 , 1 , 1 , 1 , 1 , 1 , 1 , 1 , 1 , 1 , 1 , 1 , 1 , 1 , 1 , 1 ] ) ; tf . test_util . expectArraysEqual ( await indices . data ( ) , [ 2 , 4 , 8 , 9 , 12 , 15 , 18 , 21 , 23 , 27 , 30 , 33 , 38 , 40 , 41 , 42 , 43 , 45 , 47 , 48 , 49 , 51 , 52 , 54 , 55 , 57 , 61 , 63 , 0 , 1 , 3 , 5 , 6 , 7 , 10 , 11 , 13 , 14 , 16 , 17 , 19 , 20 , 22 , 24 , 25 , 26 , 28 , 29 , 31 , 32 , 34 , 35 , 36 , 37 , 39 , 44 , 46 , 50 , 53 , 56 , 58 , 59 , 60 , 62 , 64 ] ) ; } ) ; } ) ;