Add booleanMask op (#1749)

FEATURE

Add tf.booleanMask op. Feature requested in [tensorflow/tfjs#380](tensorflow/tfjs#380).

Reference:
* [tf.boolean_mask documentation](https://www.tensorflow.org/versions/r2.0/api_docs/python/tf/boolean_mask?hl=en)
* [tf.boolean_mask tensorflow python implementation](https://github.com/tensorflow/tensorflow/blob/r2.0/tensorflow/python/ops/array_ops.py#L1274)

src/ops/boolean_mask.ts

@@ -0,0 +1,87 @@
+/**
+ * @license
+ * Copyright 2018 Google Inc. All Rights Reserved.
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * =============================================================================
+ */
+
+import {Tensor} from '../tensor';
+import {convertToTensor} from '../tensor_util_env';
+import {TensorLike} from '../types';
+import * as util from '../util';
+
+import {whereAsync} from './logical_ops';
+import {gather} from './segment_ops';
+
+/**
+ * Apply boolean mask to tensor.
+ *
+ * ```js
+ * const tensor = tf.tensor2d([1, 2, 3, 4, 5, 6], [3, 2]);
+ * const mask = tf.tensor1d([1, 0, 1], 'bool');
+ * const result = await tf.booleanMask(tensor, mask);
+ * result.print();
+ * ```
+ *
+ * @param N-D tensor.
+ * @param mask K-D boolean tensor, K <= N and K must be known statically.
+ * @param axis A 0-D int Tensor representing the axis in tensor to mask from.
+ *     By default, axis is 0 which will mask from the first dimension.
+ *     Otherwise K + axis <= N.
+ */
+/** @doc {heading: 'Tensors', subheading: 'Slicing and Joining'} */
+async function booleanMask_(
+    tensor: Tensor|TensorLike, mask: Tensor|TensorLike,
+    axis?: number): Promise<Tensor> {
+  const $tensor = convertToTensor(tensor, 'tensor', 'boolMask');
+  const $mask = convertToTensor(mask, 'mask', 'boolMask', 'bool');
+
+  const axisFrom = axis == null ? 0 : axis;
+  const maskDim = $mask.rank;
+  const tensorShape = $tensor.shape;
+
+  util.assert(maskDim > 0, () => 'mask cannot be scalar');
+  util.assertShapesMatch(
+      tensorShape.slice(axisFrom, axisFrom + maskDim), $mask.shape,
+      `mask's shape must match the first K dimensions of tensor's shape,`);
+
+  let leadingSize = 1;
+  for (let i = axisFrom; i < axisFrom + maskDim; i++) {
+    leadingSize *= tensorShape[i];
+  }
+  const targetTensorShape =
+      tensorShape.slice(0, axisFrom)
+          .concat([leadingSize], tensorShape.slice(axisFrom + maskDim));
+  const reshapedTensor = $tensor.reshape(targetTensorShape);
+  const reshapedMask = $mask.reshape([-1]);
+  const positivePositions = await whereAsync(reshapedMask);
+  const indices = positivePositions.squeeze([1]);
+
+  const res = gather(reshapedTensor, indices, axisFrom);
+
+  // Ensure no memory leak.
+  if (tensor !== $tensor) {
+    $tensor.dispose();
+  }
+  if (mask !== $mask) {
+    $mask.dispose();
+  }
+  indices.dispose();
+  reshapedTensor.dispose();
+  reshapedMask.dispose();
+  positivePositions.dispose();
+
+  return res;
+}
+
+export const booleanMask = booleanMask_;

src/ops/boolean_mask_test.ts

@@ -0,0 +1,119 @@
+/**
+ * @license
+ * Copyright 2018 Google Inc. All Rights Reserved.
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * =============================================================================
+ */
+
+import * as tf from '../index';
+import {ALL_ENVS, describeWithFlags} from '../jasmine_util';
+import {Tensor} from '../tensor';
+import {expectArraysClose} from '../test_util';
+
+describeWithFlags('booleanMask', ALL_ENVS, () => {
+  it('1d array, 1d mask, default axis', async () => {
+    const array = tf.tensor1d([1, 2, 3]);
+    const mask = tf.tensor1d([1, 0, 1], 'bool');
+    const result = await tf.booleanMask(array, mask);
+    expect(result.shape).toEqual([2]);
+    expect(result.dtype).toBe('float32');
+    expectArraysClose(await result.data(), [1, 3]);
+  });
+
+  it('2d array, 1d mask, default axis', async () => {
+    const array = tf.tensor2d([1, 2, 3, 4, 5, 6], [3, 2]);
+    const mask = tf.tensor1d([1, 0, 1], 'bool');
+    const result = await tf.booleanMask(array, mask);
+    expect(result.shape).toEqual([2, 2]);
+    expect(result.dtype).toBe('float32');
+    expectArraysClose(await result.data(), [1, 2, 5, 6]);
+  });
+
+  it('2d array, 2d mask, default axis', async () => {
+    const array = tf.tensor2d([1, 2, 3, 4, 5, 6], [3, 2]);
+    const mask = tf.tensor2d([1, 0, 1, 0, 1, 0], [3, 2], 'bool');
+    const result = await tf.booleanMask(array, mask);
+    expect(result.shape).toEqual([3]);
+    expect(result.dtype).toBe('float32');
+    expectArraysClose(await result.data(), [1, 3, 5]);
+  });
+
+  it('2d array, 1d mask, axis=1', async () => {
+    const array = tf.tensor2d([1, 2, 3, 4, 5, 6], [3, 2]);
+    const mask = tf.tensor1d([0, 1], 'bool');
+    const axis = 1;
+    const result = await tf.booleanMask(array, mask, axis);
+    expect(result.shape).toEqual([3, 1]);
+    expect(result.dtype).toBe('float32');
+    expectArraysClose(await result.data(), [2, 4, 6]);
+  });
+
+  it('accepts tensor-like object as array or mask', async () => {
+    const array = [[1, 2], [3, 4], [5, 6]];
+    const mask = [1, 0, 1];
+    const result = await tf.booleanMask(array, mask);
+    expect(result.shape).toEqual([2, 2]);
+    expect(result.dtype).toBe('float32');
+    expectArraysClose(await result.data(), [1, 2, 5, 6]);
+  });
+
+  it('ensure no memory leak', async () => {
+    const numTensorsBefore = tf.memory().numTensors;
+
+    const array = tf.tensor1d([1, 2, 3]);
+    const mask = tf.tensor1d([1, 0, 1], 'bool');
+    let resultPromise: Promise<Tensor> = null;
+
+    tf.tidy(() => {
+      resultPromise = tf.booleanMask(array, mask);
+    });
+
+    const result = await resultPromise;
+    expect(result.shape).toEqual([2]);
+    expect(result.dtype).toBe('float32');
+    expectArraysClose(await result.data(), [1, 3]);
+    array.dispose();
+    mask.dispose();
+    result.dispose();
+
+    const numTensorsAfter = tf.memory().numTensors;
+    expect(numTensorsAfter).toBe(numTensorsBefore);
+  });
+
+  it('should throw if mask is scalar', async () => {
+    const array = tf.tensor2d([1, 2, 3, 4, 5, 6], [3, 2]);
+    const mask = tf.scalar(1, 'bool');
+    let errorMessage = 'No error thrown.';
+    try {
+      await tf.booleanMask(array, mask);
+    } catch (error) {
+      errorMessage = error.message;
+    }
+    expect(errorMessage).toBe('mask cannot be scalar');
+  });
+
+  it('should throw if array and mask shape miss match', async () => {
+    const array = tf.tensor2d([1, 2, 3, 4, 5, 6], [3, 2]);
+    const mask = tf.tensor2d([1, 0], [1, 2], 'bool');
+    let errorMessage = 'No error thrown.';
+    try {
+      await tf.booleanMask(array, mask);
+    } catch (error) {
+      errorMessage = error.message;
+    }
+    expect(errorMessage)
+        .toBe(
+            `mask's shape must match the first K ` +
+            `dimensions of tensor's shape, Shapes 3,2 and 1,2 must match`);
+  });
+});

src/ops/ops.ts

@@ -16,6 +16,7 @@
  */
 
 export * from './batchnorm';
+export * from './boolean_mask';
 export * from './complex_ops';
 export * from './concat_split';
 export * from './conv';

src/tests.ts

@@ -41,6 +41,7 @@ import './ops/array_ops_test';
 import './ops/axis_util_test';
 import './ops/batchnorm_test';
 import './ops/binary_ops_test';
+import './ops/boolean_mask_test';
 import './ops/broadcast_util_test';
 import './ops/clone_test';
 import './ops/compare_ops_test';