Cleanup intopk to remove it as a kernel.

src/backends/backend.ts

@@ -241,10 +241,6 @@ export class KernelBackend implements TensorStorage, Backend, BackendTimer {
     throw new Error('Not yet implemented');
   }
 
-  inTopK<T extends Tensor, U extends Tensor>(
-      predictions: T, targets: U, k: number): U {
-    throw new Error('Not yet implemented');
-  }
 
   min(x: Tensor, axes: number[]): Tensor {
     throw new Error('Not yet implemented');

src/backends/cpu/backend_cpu.ts

@@ -41,7 +41,6 @@ import {getArrayFromDType, inferDtype, now, sizeFromShape} from '../../util';
 import {BackendTimingInfo, DataStorage, EPSILON_FLOAT32, KernelBackend} from '../backend';
 import * as backend_util from '../backend_util';
 import * as complex_util from '../complex_util';
-import {inTopKImpl} from '../inTopK_impl';
 import {nonMaxSuppressionImpl} from '../non_max_suppression_impl';
 import {split} from '../split_shared';
 import {tile} from '../tile_impl';
@@ -859,18 +858,6 @@ export class MathBackendCPU implements KernelBackend {
     return topkImpl(xVals, x.shape, x.dtype as NumericDataType, k, sorted);
   }
 
-  inTopK<T extends Tensor, U extends Tensor>(
-      predictions: T, targets: U, k: number): U {
-    this.assertNotComplex([predictions, targets], 'inTopK');
-
-    const predictionsVals = this.readSync(predictions.dataId) as TypedArray;
-    const targetsVals = this.readSync(targets.dataId) as TypedArray;
-
-    return inTopKImpl(
-               predictionsVals, predictions.shape, targetsVals, targets.shape,
-               k) as U;
-  }
-
   min(x: Tensor, axes: number[]): Tensor {
     this.assertNotComplex(x, 'min');
 

src/backends/webgl/backend_webgl.ts

@@ -44,7 +44,6 @@ import {getArrayFromDType, getTypedArrayFromDType, inferDtype, sizeFromShape} fr
 import {DataStorage, EPSILON_FLOAT16, EPSILON_FLOAT32, KernelBackend} from '../backend';
 import * as backend_util from '../backend_util';
 import {mergeRealAndImagArrays} from '../complex_util';
-import {inTopKImpl} from '../inTopK_impl';
 import {nonMaxSuppressionImpl} from '../non_max_suppression_impl';
 import {split} from '../split_shared';
 import {tile} from '../tile_impl';
@@ -1360,15 +1359,6 @@ export class MathBackendWebGL implements KernelBackend {
     return topkImpl(xVals, x.shape, x.dtype as NumericDataType, k, sorted);
   }
 
-  inTopK<T extends Tensor, U extends Tensor>(
-      predictions: T, targets: U, k: number): U {
-    const predictionsVals = predictions.dataSync();
-    const targetsVals = targets.dataSync();
-    return inTopKImpl(
-               predictionsVals, predictions.shape, targetsVals, targets.shape,
-               k) as U;
-  }
-
   min(x: Tensor, axes: number[]): Tensor {
     axis_util.assertAxesAreInnerMostDims('min', axes, x.rank);
     const [outShape, reduceShape] =

src/ops/inTopK.ts → src/ops/in_top_k.ts

@@ -15,21 +15,19 @@
  * =============================================================================
  */
 
-import {ENGINE} from '../engine';
-import {NumericTensor, Tensor} from '../tensor';
+import {Tensor} from '../tensor';
 import {convertToTensor} from '../tensor_util_env';
 import {TensorLike} from '../types';
-import {assert, assertShapesMatch} from '../util';
-
-import {op} from './operation';
+import {assert, assertShapesMatch, getTypedArrayFromDType} from '../util';
+import {tensor} from './tensor_ops';
 
 /**
  * Says whether the targets are in the top K predictions.
  *
  * ```js
  * const predictions = tf.tensor2d([[20, 10, 40, 30], [30, 50, -20, 10]]);
  * const targets = tf.tensor1d([2, 0]);
- * const precision = tf.inTopK(predictions, targets);
+ * const precision = await tf.inTopKAsync(predictions, targets);
  * precision.print();
  * ```
  * @param predictions 2-D or higher `tf.Tensor` with last dimension being
@@ -39,8 +37,8 @@ import {op} from './operation';
  *     default to 1.
  */
 /** @doc {heading: 'Operations', subheading: 'Evaluation'} */
-function inTopK_<T extends Tensor, U extends Tensor>(
-    predictions: T|TensorLike, targets: U|TensorLike, k = 1): U {
+async function inTopKAsync_<T extends Tensor, U extends Tensor>(
+    predictions: T|TensorLike, targets: U|TensorLike, k = 1): Promise<U> {
   const $predictions = convertToTensor(predictions, 'predictions', 'inTopK');
   const $targets = convertToTensor(targets, 'targets', 'inTopK');
 
@@ -50,9 +48,9 @@ function inTopK_<T extends Tensor, U extends Tensor>(
           `but got ${$predictions.rank}`);
   assert(
       $predictions.rank - 1 === $targets.rank,
-      () => `predictions' rank should be 1 larger than ` +
-          `targets' rank, but got predictions' rank ` +
-          `${$predictions.rank} and targets' rank ${$targets.rank}`);
+      () => `predictions rank should be 1 larger than ` +
+          `targets rank, but got predictions rank ` +
+          `${$predictions.rank} and targets rank ${$targets.rank}`);
   assertShapesMatch(
       $predictions.shape.slice(0, $predictions.shape.length - 1),
       $targets.shape,
@@ -61,15 +59,44 @@ function inTopK_<T extends Tensor, U extends Tensor>(
   const lastDim = $predictions.shape[$predictions.shape.length - 1];
   assert(
       k > 0 && k <= lastDim,
-      () => `'k' passed to inTopK() must be > 0 && <= the predictions' last ` +
+      () => `'k' passed to inTopK() must be > 0 && <= the predictions last ` +
           `dimension (${lastDim}), but got ${k}`);
 
-  const precision = ENGINE.runKernel(
-      b =>
-          b.inTopK($predictions as NumericTensor, $targets as NumericTensor, k),
-      {$predictions, $targets});
+  const predictionsVals = await $predictions.data();
+  const targetsVals = await $targets.data();
+
+  // Reshape predictionsVals into a 2d tensor [batch, lastDim]
+  // and look up topK along lastDim.
+  const [batch, size] = [predictionsVals.length / lastDim, lastDim];
+  const precision = getTypedArrayFromDType('bool', batch);
+
+  for (let b = 0; b < batch; b++) {
+    const offset = b * size;
+    const vals = predictionsVals.subarray(offset, offset + size);
+    const valAndInd: Array<{value: number, index: number}> = [];
+    for (let i = 0; i < vals.length; i++) {
+      valAndInd.push({value: vals[i], index: i});
+    }
+    valAndInd.sort((a, b) => b.value - a.value);
+
+    precision[b] = 0;
+    for (let i = 0; i < k; i++) {
+      if (valAndInd[i].index === targetsVals[b]) {
+        precision[b] = 1;
+        break;
+      }
+    }
+  }
+
+  if (predictions !== $predictions) {
+    $predictions.dispose();
+  }
+  if (targets !== $targets) {
+    $targets.dispose();
+  }
 
-  return precision as U;
+  // Output precision has the same shape as targets.
+  return tensor(precision, $targets.shape, 'bool') as U;
 }
 
-export const inTopK = op({inTopK_});
+export const inTopKAsync = inTopKAsync_;

src/ops/inTopK_test.ts → src/ops/in_top_k_test.ts

@@ -21,11 +21,11 @@ import {expectArraysClose} from '../test_util';
 
 import {tensor1d, tensor2d, tensor3d} from './tensor_ops';
 
-describeWithFlags('inTopK', ALL_ENVS, async () => {
+describeWithFlags('inTopKAsync', ALL_ENVS, async () => {
   it('predictions 2d array, targets 1d array, with default k', async () => {
     const predictions = tensor2d([[20, 10, 40, 30], [30, 50, -20, 10]]);
     const targets = tensor1d([2, 0]);
-    const precision = tf.inTopK(predictions, targets);
+    const precision = await tf.inTopKAsync(predictions, targets);
     expect(precision.shape).toEqual([2]);
     expect(precision.dtype).toBe('bool');
     expectArraysClose(await precision.data(), [1, 0]);
@@ -35,7 +35,7 @@ describeWithFlags('inTopK', ALL_ENVS, async () => {
     const predictions = tensor2d([[20, 10, 40, 30], [30, 50, -20, 10]]);
     const targets = tensor1d([2, 0]);
     const k = 2;
-    const precision = tf.inTopK(predictions, targets, k);
+    const precision = await tf.inTopKAsync(predictions, targets, k);
     expect(precision.shape).toEqual([2]);
     expect(precision.dtype).toBe('bool');
     expectArraysClose(await precision.data(), [1, 1]);
@@ -45,7 +45,7 @@ describeWithFlags('inTopK', ALL_ENVS, async () => {
     const predictions =
         tensor3d([[[1, 5, 2], [4, 3, 6]], [[3, 2, 1], [1, 2, 3]]]);
     const targets = tensor2d([[1, 2], [0, 1]]);
-    const precision = tf.inTopK(predictions, targets);
+    const precision = await tf.inTopKAsync(predictions, targets);
     expect(precision.shape).toEqual([2, 2]);
     expect(precision.dtype).toBe('bool');
     expectArraysClose(await precision.data(), [1, 1, 1, 0]);
@@ -56,7 +56,7 @@ describeWithFlags('inTopK', ALL_ENVS, async () => {
         tensor3d([[[1, 5, 2], [4, 3, 6]], [[3, 2, 1], [1, 2, 3]]]);
     const targets = tensor2d([[1, 2], [0, 1]]);
     const k = 2;
-    const precision = tf.inTopK(predictions, targets, k);
+    const precision = await tf.inTopKAsync(predictions, targets, k);
     expect(precision.shape).toEqual([2, 2]);
     expect(precision.dtype).toBe('bool');
     expectArraysClose(await precision.data(), [1, 1, 1, 1]);
@@ -66,13 +66,13 @@ describeWithFlags('inTopK', ALL_ENVS, async () => {
     const predictions = tensor2d([[1, 2, 2, 1]]);
 
     const targets1 = tensor1d([1]);
-    const precision1 = tf.inTopK(predictions, targets1);
+    const precision1 = await tf.inTopKAsync(predictions, targets1);
     expect(precision1.shape).toEqual([1]);
     expect(precision1.dtype).toBe('bool');
     expectArraysClose(await precision1.data(), [1]);
 
     const targets2 = tensor1d([2]);
-    const precision2 = tf.inTopK(predictions, targets2);
+    const precision2 = await tf.inTopKAsync(predictions, targets2);
     expect(precision2.shape).toEqual([1]);
     expect(precision2.dtype).toBe('bool');
     expectArraysClose(await precision2.data(), [0]);
@@ -81,28 +81,72 @@ describeWithFlags('inTopK', ALL_ENVS, async () => {
   it('accept tensor-like object, with default k', async () => {
     const predictions = [[20, 10, 40, 30], [30, 50, -20, 10]];
     const targets = [2, 0];
-    const precision = tf.inTopK(predictions, targets);
+    const precision = await tf.inTopKAsync(predictions, targets);
     expect(precision.shape).toEqual([2]);
     expect(precision.dtype).toBe('bool');
     expectArraysClose(await precision.data(), [1, 0]);
   });
 
-  it('throws when predictions_rank <2', () => {
+  it('doesnt leak tensors with tensor-like objects', async () => {
+    const numTensors = tf.memory().numTensors;
+
+    const predictions = [[20, 10, 40, 30], [30, 50, -20, 10]];
+    const targets = [2, 0];
+    const precision = await tf.inTopKAsync(predictions, targets);
+    precision.dispose();
+
+    expect(tf.memory().numTensors).toBe(numTensors);
+  });
+
+  it('throws when predictions_rank <2', async () => {
     const predictions = tensor1d([20, 10, 40, 30]);
     const targets = [2];
-    expect(() => tf.inTopK(predictions, targets)).toThrowError();
+
+    // expect(...).toThrowError() does not support async functions.
+    // See https://github.com/jasmine/jasmine/issues/1410
+    try {
+      await tf.inTopKAsync(predictions, targets);
+      throw new Error('The line above should have thrown an error');
+    } catch (ex) {
+      expect(ex.message)
+          .toEqual(
+              'inTopK() expects the predictions to ' +
+              'be of rank 2 or higher, but got 1');
+    }
   });
 
-  it('throws when prediction_rank != targets_rank + 1', () => {
+  it('throws when prediction.rank != targets.rank + 1', async () => {
     const predictions = tensor2d([[20, 10, 40, 30], [30, 50, -20, 10]]);
     const targets = tensor2d([[0], [0]]);
-    expect(() => tf.inTopK(predictions, targets)).toThrowError();
+
+    // expect(...).toThrowError() does not support async functions.
+    // See https://github.com/jasmine/jasmine/issues/1410
+    try {
+      await tf.inTopKAsync(predictions, targets);
+      throw new Error('The line above should have thrown an error');
+    } catch (ex) {
+      expect(ex.message)
+          .toEqual(
+              'predictions rank should be 1 larger than targets rank,' +
+              ' but got predictions rank 2 and targets rank 2');
+    }
   });
 
-  it('throws when k > size of last dimension of predictions', () => {
+  it('throws when k > size of last dimension of predictions', async () => {
     const predictions = tensor2d([[20, 10, 40, 30], [30, 50, -20, 10]]);
     const targets = tensor1d([2, 0]);
     const k = 5;
-    expect(() => tf.inTopK(predictions, targets, k)).toThrowError();
+
+    // expect(...).toThrowError() does not support async functions.
+    // See https://github.com/jasmine/jasmine/issues/1410
+    try {
+      await tf.inTopKAsync(predictions, targets, k);
+      throw new Error('The line above should have thrown an error');
+    } catch (ex) {
+      expect(ex.message)
+          .toEqual(
+              '\'k\' passed to inTopK() must be > 0 && <= the predictions ' +
+              'last dimension (4), but got 5');
+    }
   });
 });

src/ops/ops.ts

@@ -48,7 +48,7 @@ export * from './gather_nd';
 export * from './diag';
 export * from './dropout';
 export * from './signal_ops';
-export * from './inTopK';
+export * from './in_top_k';
 
 export {op} from './operation';
 

src/tests.ts

@@ -61,7 +61,7 @@ import './ops/dropout_test';
 import './ops/fused_test';
 import './ops/gather_nd_test';
 import './ops/image_ops_test';
-import './ops/inTopK_test';
+import './ops/in_top_k_test';
 import './ops/linalg_ops_test';
 import './ops/logical_ops_test';
 import './ops/loss_ops_test';