| 1 | {"version":3,"file":"topk_test.js","sourceRoot":"","sources":["../../src/ops/topk_test.ts"],"names":[],"mappings":"AAAA;;;;;;;;;;;;;;;GAeG;AAEH,OAAO,KAAK,EAAE,MAAM,UAAU,CAAC;AAC/B,OAAO,EAAC,QAAQ,EAAE,iBAAiB,EAAC,MAAM,iBAAiB,CAAC;AAC5D,OAAO,EAAC,iBAAiB,EAAC,MAAM,cAAc,CAAC;AAE/C,OAAO,EAAC,MAAM,EAAC,MAAM,UAAU,CAAC;AAChC,OAAO,EAAC,QAAQ,EAAC,MAAM,YAAY,CAAC;AACpC,OAAO,EAAC,QAAQ,EAAC,MAAM,YAAY,CAAC;AACpC,OAAO,EAAC,QAAQ,EAAC,MAAM,YAAY,CAAC;AAEpC,iBAAiB,CAAC,MAAM,EAAE,QAAQ,EAAE,GAAG,EAAE;IACvC,EAAE,CAAC,yBAAyB,EAAE,KAAK,IAAI,EAAE;QACvC,MAAM,CAAC,GAAG,QAAQ,CAAC,CAAC,EAAE,EAAE,EAAE,EAAE,EAAE,EAAE,EAAE,CAAC,CAAC,CAAC;QACrC,MAAM,EAAC,MAAM,EAAE,OAAO,EAAC,GAAG,EAAE,CAAC,IAAI,CAAC,CAAC,CAAC,CAAC;QAErC,MAAM,CAAC,MAAM,CAAC,KAAK,CAAC,CAAC,OAAO,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC;QAClC,MAAM,CAAC,OAAO,CAAC,KAAK,CAAC,CAAC,OAAO,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC;QACnC,MAAM,CAAC,MAAM,CAAC,KAAK,CAAC,CAAC,IAAI,CAAC,SAAS,CAAC,CAAC;QACrC,MAAM,CAAC,OAAO,CAAC,KAAK,CAAC,CAAC,IAAI,CAAC,OAAO,CAAC,CAAC;QACpC,iBAAiB,CAAC,MAAM,MAAM,CAAC,IAAI,EAAE,EAAE,CAAC,EAAE,CAAC,CAAC,CAAC;QAC7C,iBAAiB,CAAC,MAAM,OAAO,CAAC,IAAI,EAAE,EAAE,CAAC,CAAC,CAAC,CAAC,CAAC;IAC/C,CAAC,CAAC,CAAC;IAEH,EAAE,CAAC,0CAA0C,EAAE,KAAK,IAAI,EAAE;QACxD,MAAM,CAAC,GAAG,QAAQ,CAAC,CAAC,EAAE,EAAE,EAAE,EAAE,EAAE,EAAE,EAAE,CAAC,CAAC,CAAC;QACrC,MAAM,EAAC,MAAM,EAAE,OAAO,EAAC,GAAG,CAAC,CAAC,IAAI,EAAE,CAAC;QAEnC,MAAM,CAAC,MAAM,CAAC,KAAK,CAAC,CAAC,OAAO,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC;QAClC,MAAM,CAAC,OAAO,CAAC,KAAK,CAAC,CAAC,OAAO,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC;QACnC,MAAM,CAAC,MAAM,CAAC,KAAK,CAAC,CAAC,IAAI,CAAC,SAAS,CAAC,CAAC;QACrC,MAAM,CAAC,OAAO,CAAC,KAAK,CAAC,CAAC,IAAI,CAAC,OAAO,CAAC,CAAC;QACpC,iBAAiB,CAAC,MAAM,MAAM,CAAC,IAAI,EAAE,EAAE,CAAC,EAAE,CAAC,CAAC,CAAC;QAC7C,iBAAiB,CAAC,MAAM,OAAO,CAAC,IAAI,EAAE,EAAE,CAAC,CAAC,CAAC,CAAC,CAAC;IAC/C,CAAC,CAAC,CAAC;IAEH,EAAE,CAAC,yBAAyB,EAAE,KAAK,IAAI,EAAE;QACvC,MAAM,CAAC,GAAG,QAAQ,CAAC,CAAC,CAAC,EAAE,EAAE,EAAE,CAAC,EAAE,CAAC,EAAE,EAAE,EAAE,CAAC,CAAC,CAAC,CAAC;QACzC,MAAM,EAAC,MAAM,EAAE,OAAO,EAAC,GAAG,EAAE,CAAC,IAAI,CAAC,CAAC,CAAC,CAAC;QAErC,MAAM,CAAC,MAAM,CAAC,KAAK,CAAC,CAAC,OAAO,CAAC,CAAC,CAAC,EAAE,CAAC,CAAC,CAAC,CAAC;QACrC,MAAM,CAAC,OAAO,CAAC,KAAK,CAAC,CAAC,OAAO,CAAC,CAAC,CAAC,EAAE,CAAC,CAAC,CAAC,CAAC;QACtC,MAAM,CAAC,MAAM,CAAC,KAAK,CAAC,CAAC,IAAI,CAAC,SAAS,CAAC,CAAC;QACrC,MAAM,CAAC,OAAO,CAAC,KAAK,CAAC,CAAC,IAAI,CAAC,OAAO,CAAC,CAAC;QACpC,iBAAiB,CAAC,MAAM,MAAM,CAAC,IAAI,EAAE,EAAE,CAAC,EAAE,EAAE,EAAE,CAAC,CAAC,CAAC;QACjD,iBAAiB,CAAC,MAAM,OAAO,CAAC,IAAI,EAAE,EAAE,CAAC,CAAC,EAAE,CAAC,CAAC,CAAC,CAAC;IAClD,CAAC,CAAC,CAAC;IAEH,EAAE,CAAC,mBAAmB,EAAE,KAAK,IAAI,EAAE;QACjC,MAAM,CAAC,GAAG,QAAQ,CAAC;YACjB,CAAC,CAAC,EAAE,CAAC,EAAE,CAAC,CAAC;YACT,CAAC,CAAC,EAAE,CAAC,EAAE,CAAC,CAAC;YACT,CAAC,CAAC,EAAE,CAAC,EAAE,CAAC,CAAC;YACT,CAAC,CAAC,EAAE,CAAC,EAAE,CAAC,CAAC;SACV,CAAC,CAAC;QACH,MAAM,CAAC,GAAG,CAAC,CAAC;QACZ,MAAM,EAAC,MAAM,EAAE,OAAO,EAAC,GAAG,EAAE,CAAC,IAAI,CAAC,CAAC,EAAE,CAAC,CAAC,CAAC;QAExC,MAAM,CAAC,MAAM,CAAC,KAAK,CAAC,CAAC,OAAO,CAAC,CAAC,CAAC,EAAE,CAAC,CAAC,CAAC,CAAC;QACrC,MAAM,CAAC,OAAO,CAAC,KAAK,CAAC,CAAC,OAAO,CAAC,CAAC,CAAC,EAAE,CAAC,CAAC,CAAC,CAAC;QACtC,MAAM,CAAC,MAAM,CAAC,KAAK,CAAC,CAAC,IAAI,CAAC,SAAS,CAAC,CAAC;QACrC,MAAM,CAAC,OAAO,CAAC,KAAK,CAAC,CAAC,IAAI,CAAC,OAAO,CAAC,CAAC;QACpC,iBAAiB,CAAC,MAAM,MAAM,CAAC,IAAI,EAAE,EAAE,CAAC,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,CAAC,CAAC,CAAC;QACjE,iBAAiB,CAAC,MAAM,OAAO,CAAC,IAAI,EAAE,EAAE,CAAC,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,CAAC,CAAC,CAAC;IACpE,CAAC,CAAC,CAAC;IAEH,EAAE,CAAC,oCAAoC,EAAE,KAAK,IAAI,EAAE;QAClD,MAAM,CAAC,GAAG,QAAQ,CAAC;YACjB,CAAC,CAAC,EAAE,CAAC,EAAE,CAAC,CAAC;YACT,CAAC,CAAC,EAAE,CAAC,EAAE,CAAC,CAAC;YACT,CAAC,CAAC,EAAE,CAAC,EAAE,CAAC,CAAC;YACT,CAAC,CAAC,EAAE,CAAC,EAAE,CAAC,CAAC;SACV,CAAC,CAAC;QACH,MAAM,CAAC,GAAG,CAAC,CAAC;QACZ,MAAM,EAAC,MAAM,EAAE,OAAO,EAAC,GAAG,CAAC,CAAC,IAAI,CAAC,CAAC,CAAC,CAAC;QAEpC,MAAM,CAAC,MAAM,CAAC,KAAK,CAAC,CAAC,OAAO,CAAC,CAAC,CAAC,EAAE,CAAC,CAAC,CAAC,CAAC;QACrC,MAAM,CAAC,OAAO,CAAC,KAAK,CAAC,CAAC,OAAO,CAAC,CAAC,CAAC,EAAE,CAAC,CAAC,CAAC,CAAC;QACtC,MAAM,CAAC,MAAM,CAAC,KAAK,CAAC,CAAC,IAAI,CAAC,SAAS,CAAC,CAAC;QACrC,MAAM,CAAC,OAAO,CAAC,KAAK,CAAC,CAAC,IAAI,CAAC,OAAO,CAAC,CAAC;QACpC,iBAAiB,CAAC,MAAM,MAAM,CAAC,IAAI,EAAE,EAAE,CAAC,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,CAAC,CAAC,CAAC;QACjE,iBAAiB,CAAC,MAAM,OAAO,CAAC,IAAI,EAAE,EAAE,CAAC,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,CAAC,CAAC,CAAC;IACpE,CAAC,CAAC,CAAC;IAEH,EAAE,CAAC,mBAAmB,EAAE,KAAK,IAAI,EAAE;QACjC,MAAM,CAAC,GAAG,QAAQ,CAAC;YACjB,CAAC,CAAC,CAAC,EAAE,CAAC,EAAE,CAAC,CAAC,EAAE,CAAC,CAAC,EAAE,CAAC,EAAE,CAAC,CAAC,CAAC;YACtB,CAAC,CAAC,CAAC,EAAE,CAAC,EAAE,CAAC,CAAC,EAAE,CAAC,CAAC,EAAE,CAAC,EAAE,CAAC,CAAC,CAAC;SACvB,CAAC,CAAC,CAAE,SAAS;QACd,MAAM,CAAC,GAAG,CAAC,CAAC;QACZ,MAAM,EAAC,MAAM,EAAE,OAAO,EAAC,GAAG,EAAE,CAAC,IAAI,CAAC,CAAC,EAAE,CAAC,CAAC,CAAC;QAExC,MAAM,CAAC,MAAM,CAAC,KAAK,CAAC,CAAC,OAAO,CAAC,CAAC,CAAC,EAAE,CAAC,EAAE,CAAC,CAAC,CAAC,CAAC;QACxC,MAAM,CAAC,OAAO,CAAC,KAAK,CAAC,CAAC,OAAO,CAAC,CAAC,CAAC,EAAE,CAAC,EAAE,CAAC,CAAC,CAAC,CAAC;QACzC,MAAM,CAAC,MAAM,CAAC,KAAK,CAAC,CAAC,IAAI,CAAC,SAAS,CAAC,CAAC;QACrC,MAAM,CAAC,OAAO,CAAC,KAAK,CAAC,CAAC,IAAI,CAAC,OAAO,CAAC,CAAC;QACpC,iBAAiB,CACb,MAAM,MAAM,CAAC,IAAI,EAAE,EAAE,CAAC,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,CAAC,CAAC,CAAC;QAC/D,iBAAiB,CACb,MAAM,OAAO,CAAC,IAAI,EAAE,EAAE,CAAC,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,CAAC,CAAC,CAAC;IAClE,CAAC,CAAC,CAAC;IAEH,EAAE,CAAC,oCAAoC,EAAE,KAAK,IAAI,EAAE;QAClD,MAAM,CAAC,GAAG,QAAQ,CAAC,CAAC,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,CAAC,EAAE,OAAO,CAAC,CAAC;QAC1C,MAAM,EAAC,MAAM,EAAE,OAAO,EAAC,GAAG,EAAE,CAAC,IAAI,CAAC,CAAC,CAAC,CAAC;QAErC,MAAM,CAAC,MAAM,CAAC,KAAK,CAAC,CAAC,OAAO,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC;QAClC,MAAM,CAAC,OAAO,CAAC,KAAK,CAAC,CAAC,OAAO,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC;QACnC,MAAM,CAAC,MAAM,CAAC,KAAK,CAAC,CAAC,IAAI,CAAC,OAAO,CAAC,CAAC;QACnC,MAAM,CAAC,OAAO,CAAC,KAAK,CAAC,CAAC,IAAI,CAAC,OAAO,CAAC,CAAC;QACpC,iBAAiB,CAAC,MAAM,MAAM,CAAC,IAAI,EAAE,EAAE,CAAC,CAAC,CAAC,CAAC,CAAC;QAC5C,iBAAiB,CAAC,MAAM,OAAO,CAAC,IAAI,EAAE,EAAE,CAAC,CAAC,CAAC,CAAC,CAAC;IAC/C,CAAC,CAAC,CAAC;IAEH,EAAE,CAAC,wCAAwC,EAAE,KAAK,IAAI,EAAE;QACtD,MAAM,CAAC,GAAG,QAAQ,CAAC,CAAC,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,CAAC,EAAE,OAAO,CAAC,CAAC;QAC1C,MAAM,CAAC,GAAG,CAAC,CAAC;QACZ,MAAM,EAAC,MAAM,EAAE,OAAO,EAAC,GAAG,EAAE,CAAC,IAAI,CAAC,CAAC,EAAE,CAAC,CAAC,CAAC;QAExC,MAAM,CAAC,MAAM,CAAC,KAAK,CAAC,CAAC,OAAO,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC;QAClC,MAAM,CAAC,OAAO,CAAC,KAAK,CAAC,CAAC,OAAO,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC;QACnC,MAAM,CAAC,MAAM,CAAC,KAAK,CAAC,CAAC,IAAI,CAAC,OAAO,CAAC,CAAC;QACnC,MAAM,CAAC,OAAO,CAAC,KAAK,CAAC,CAAC,IAAI,CAAC,OAAO,CAAC,CAAC;QACpC,iBAAiB,CAAC,MAAM,MAAM,CAAC,IAAI,EAAE,EAAE,CAAC,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,CAAC,CAAC,CAAC;QACrD,iBAAiB,CAAC,MAAM,OAAO,CAAC,IAAI,EAAE,EAAE,CAAC,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,CAAC,CAAC,CAAC;IACxD,CAAC,CAAC,CAAC;IAEH,EAAE,CAAC,+BAA+B,EAAE,GAAG,EAAE;QACvC,MAAM,CAAC,GAAG,QAAQ,CAAC,CAAC,CAAC,EAAE,EAAE,EAAE,CAAC,EAAE,CAAC,EAAE,EAAE,EAAE,CAAC,CAAC,CAAC,CAAC;QACzC,MAAM,CAAC,GAAG,EAAE,CAAC,EAAE,CAAC,IAAI,CAAC,CAAC,EAAE,CAAC,CAAC,CAAC;aACtB,YAAY,CAAC,sDAAsD,CAAC,CAAC;IAC5E,CAAC,CAAC,CAAC;IAEH,EAAE,CAAC,6BAA6B,EAAE,GAAG,EAAE;QACrC,MAAM,CAAC,GAAG,MAAM,CAAC,CAAC,CAAC,CAAC;QACpB,MAAM,CAAC,GAAG,EAAE,CAAC,EAAE,CAAC,IAAI,CAAC,CAAC,CAAC,CAAC;aACnB,YAAY,CAAC,sDAAsD,CAAC,CAAC;IAC5E,CAAC,CAAC,CAAC;IAEH,EAAE,CAAC,mCAAmC,EAAE,KAAK,IAAI,EAAE;QACjD,MAAM,CAAC,GAAG,CAAC,EAAE,EAAE,EAAE,EAAE,EAAE,EAAE,EAAE,CAAC,CAAC;QAC3B,MAAM,CAAC,GAAG,CAAC,CAAC;QACZ,MAAM,EAAC,MAAM,EAAE,OAAO,EAAC,GAAG,EAAE,CAAC,IAAI,CAAC,CAAC,EAAE,CAAC,CAAC,CAAC;QAExC,MAAM,CAAC,MAAM,CAAC,KAAK,CAAC,CAAC,OAAO,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC;QAClC,MAAM,CAAC,OAAO,CAAC,KAAK,CAAC,CAAC,OAAO,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC;QACnC,MAAM,CAAC,MAAM,CAAC,KAAK,CAAC,CAAC,IAAI,CAAC,SAAS,CAAC,CAAC;QACrC,MAAM,CAAC,OAAO,CAAC,KAAK,CAAC,CAAC,IAAI,CAAC,OAAO,CAAC,CAAC;QACpC,iBAAiB,CAAC,MAAM,MAAM,CAAC,IAAI,EAAE,EAAE,CAAC,EAAE,EAAE,EAAE,CAAC,CAAC,CAAC;QACjD,iBAAiB,CAAC,MAAM,OAAO,CAAC,IAAI,EAAE,EAAE,CAAC,CAAC,EAAE,CAAC,CAAC,CAAC,CAAC;IAClD,CAAC,CAAC,CAAC;IAEH,EAAE,CAAC,uCAAuC,EAAE,KAAK,IAAI,EAAE;QACrD,MAAM,CAAC,GAAG,QAAQ,CAAC,CAAC,EAAE,EAAE,EAAE,EAAE,EAAE,EAAE,EAAE,CAAC,CAAC,CAAC;QACrC,MAAM,CAAC,GAAG,MAAM,CAAC,CAAC,CAAC,CAAC;QACpB,MAAM,EAAC,MAAM,EAAE,OAAO,EAAC,GAAG,EAAE,CAAC,IAAI,CAAC,EAAE,CAAC,GAAG,CAAC,CAAC,EAAE,CAAC,CAAC,CAAC,CAAC;QAEhD,MAAM,CAAC,MAAM,CAAC,KAAK,CAAC,CAAC,OAAO,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC;QAClC,MAAM,CAAC,OAAO,CAAC,KAAK,CAAC,CAAC,OAAO,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC;QACnC,MAAM,CAAC,MAAM,CAAC,KAAK,CAAC,CAAC,IAAI,CAAC,SAAS,CAAC,CAAC;QACrC,MAAM,CAAC,OAAO,CAAC,KAAK,CAAC,CAAC,IAAI,CAAC,OAAO,CAAC,CAAC;QACpC,iBAAiB,CAAC,MAAM,MAAM,CAAC,IAAI,EAAE,EAAE,CAAC,EAAE,CAAC,CAAC,CAAC;QAC7C,iBAAiB,CAAC,MAAM,OAAO,CAAC,IAAI,EAAE,EAAE,CAAC,CAAC,CAAC,CAAC,CAAC;IAC/C,CAAC,CAAC,CAAC;AACL,CAAC,CAAC,CAAC","sourcesContent":["/**\n * @license\n * Copyright 2018 Google LLC. All Rights Reserved.\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n * http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n * =============================================================================\n */\n\nimport * as tf from '../index';\nimport {ALL_ENVS, describeWithFlags} from '../jasmine_util';\nimport {expectArraysClose} from '../test_util';\n\nimport {scalar} from './scalar';\nimport {tensor1d} from './tensor1d';\nimport {tensor2d} from './tensor2d';\nimport {tensor3d} from './tensor3d';\n\ndescribeWithFlags('topk', ALL_ENVS, () => {\n it('1d array with default k', async () => {\n const a = tensor1d([20, 10, 40, 30]);\n const {values, indices} = tf.topk(a);\n\n expect(values.shape).toEqual([1]);\n expect(indices.shape).toEqual([1]);\n expect(values.dtype).toBe('float32');\n expect(indices.dtype).toBe('int32');\n expectArraysClose(await values.data(), [40]);\n expectArraysClose(await indices.data(), [2]);\n });\n\n it('1d array with default k from tensor.topk', async () => {\n const a = tensor1d([20, 10, 40, 30]);\n const {values, indices} = a.topk();\n\n expect(values.shape).toEqual([1]);\n expect(indices.shape).toEqual([1]);\n expect(values.dtype).toBe('float32');\n expect(indices.dtype).toBe('int32');\n expectArraysClose(await values.data(), [40]);\n expectArraysClose(await indices.data(), [2]);\n });\n\n it('2d array with default k', async () => {\n const a = tensor2d([[10, 50], [40, 30]]);\n const {values, indices} = tf.topk(a);\n\n expect(values.shape).toEqual([2, 1]);\n expect(indices.shape).toEqual([2, 1]);\n expect(values.dtype).toBe('float32');\n expect(indices.dtype).toBe('int32');\n expectArraysClose(await values.data(), [50, 40]);\n expectArraysClose(await indices.data(), [1, 0]);\n });\n\n it('2d array with k=2', async () => {\n const a = tensor2d([\n [1, 5, 2],\n [4, 3, 6],\n [3, 2, 1],\n [1, 2, 3],\n ]);\n const k = 2;\n const {values, indices} = tf.topk(a, k);\n\n expect(values.shape).toEqual([4, 2]);\n expect(indices.shape).toEqual([4, 2]);\n expect(values.dtype).toBe('float32');\n expect(indices.dtype).toBe('int32');\n expectArraysClose(await values.data(), [5, 2, 6, 4, 3, 2, 3, 2]);\n expectArraysClose(await indices.data(), [1, 2, 2, 0, 0, 1, 2, 1]);\n });\n\n it('2d array with k=2 from tensor.topk', async () => {\n const a = tensor2d([\n [1, 5, 2],\n [4, 3, 6],\n [3, 2, 1],\n [1, 2, 3],\n ]);\n const k = 2;\n const {values, indices} = a.topk(k);\n\n expect(values.shape).toEqual([4, 2]);\n expect(indices.shape).toEqual([4, 2]);\n expect(values.dtype).toBe('float32');\n expect(indices.dtype).toBe('int32');\n expectArraysClose(await values.data(), [5, 2, 6, 4, 3, 2, 3, 2]);\n expectArraysClose(await indices.data(), [1, 2, 2, 0, 0, 1, 2, 1]);\n });\n\n it('3d array with k=3', async () => {\n const a = tensor3d([\n [[1, 5, 2], [4, 3, 6]],\n [[3, 2, 1], [1, 2, 3]],\n ]); // 2x2x3.\n const k = 3;\n const {values, indices} = tf.topk(a, k);\n\n expect(values.shape).toEqual([2, 2, 3]);\n expect(indices.shape).toEqual([2, 2, 3]);\n expect(values.dtype).toBe('float32');\n expect(indices.dtype).toBe('int32');\n expectArraysClose(\n await values.data(), [5, 2, 1, 6, 4, 3, 3, 2, 1, 3, 2, 1]);\n expectArraysClose(\n await indices.data(), [1, 2, 0, 2, 0, 1, 0, 1, 2, 2, 1, 0]);\n });\n\n it('topk(int32) propagates int32 dtype', async () => {\n const a = tensor1d([2, 3, 1, 4], 'int32');\n const {values, indices} = tf.topk(a);\n\n expect(values.shape).toEqual([1]);\n expect(indices.shape).toEqual([1]);\n expect(values.dtype).toBe('int32');\n expect(indices.dtype).toBe('int32');\n expectArraysClose(await values.data(), [4]);\n expectArraysClose(await indices.data(), [3]);\n });\n\n it('lower-index element appears first, k=4', async () => {\n const a = tensor1d([1, 2, 2, 1], 'int32');\n const k = 4;\n const {values, indices} = tf.topk(a, k);\n\n expect(values.shape).toEqual([4]);\n expect(indices.shape).toEqual([4]);\n expect(values.dtype).toBe('int32');\n expect(indices.dtype).toBe('int32');\n expectArraysClose(await values.data(), [2, 2, 1, 1]);\n expectArraysClose(await indices.data(), [1, 2, 0, 3]);\n });\n\n it('throws when k > size of array', () => {\n const a = tensor2d([[10, 50], [40, 30]]);\n expect(() => tf.topk(a, 3))\n .toThrowError(/'k' passed to topk\\(\\) must be <= the last dimension/);\n });\n\n it('throws when passed a scalar', () => {\n const a = scalar(2);\n expect(() => tf.topk(a))\n .toThrowError(/topk\\(\\) expects the input to be of rank 1 or higher/);\n });\n\n it('accepts a tensor-like object, k=2', async () => {\n const a = [20, 10, 40, 30];\n const k = 2;\n const {values, indices} = tf.topk(a, k);\n\n expect(values.shape).toEqual([2]);\n expect(indices.shape).toEqual([2]);\n expect(values.dtype).toBe('float32');\n expect(indices.dtype).toBe('int32');\n expectArraysClose(await values.data(), [40, 30]);\n expectArraysClose(await indices.data(), [2, 3]);\n });\n\n it('handles output tensors from other ops', async () => {\n const a = tensor1d([20, 10, 40, 30]);\n const b = scalar(2);\n const {values, indices} = tf.topk(tf.add(a, b));\n\n expect(values.shape).toEqual([1]);\n expect(indices.shape).toEqual([1]);\n expect(values.dtype).toBe('float32');\n expect(indices.dtype).toBe('int32');\n expectArraysClose(await values.data(), [42]);\n expectArraysClose(await indices.data(), [2]);\n });\n});\n"]} |