Modularise logical ops

tfjs-core/src/gradients/Pow_grad.ts

@@ -19,13 +19,13 @@ import {GradConfig} from '../kernel_registry';
 import {cast, reshape} from '../ops/array_ops';
 import * as broadcast_util from '../ops/broadcast_util';
 import {greater} from '../ops/greater';
-import {where} from '../ops/logical_ops';
 import {mul} from '../ops/mul';
 import {pow} from '../ops/pow';
 import {sum} from '../ops/reduction_ops';
 import {sub} from '../ops/sub';
 import {scalar, zerosLike} from '../ops/tensor_ops';
 import {log} from '../ops/unary_ops';
+import {where} from '../ops/where';
 import {Tensor} from '../tensor';
 
 export const powGradConfig: GradConfig = {

tfjs-core/src/gradients/Prelu_grad.ts

@@ -19,10 +19,10 @@ import {GradConfig} from '../kernel_registry';
 import {reshape} from '../ops/array_ops';
 import {getReductionAxes} from '../ops/broadcast_util';
 import {greater} from '../ops/greater';
-import {where} from '../ops/logical_ops';
 import {mul} from '../ops/mul';
 import {sum} from '../ops/reduction_ops';
 import {zerosLike} from '../ops/tensor_ops';
+import {where} from '../ops/where';
 import {Tensor} from '../tensor';
 
 export const preluGradConfig: GradConfig = {

tfjs-core/src/gradients/SelectV2_grad.ts

@@ -0,0 +1,39 @@
+/**
+ * @license
+ * Copyright 2020 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 {SelectV2} from '../kernel_names';
+import {GradConfig} from '../kernel_registry';
+import {cast} from '../ops/array_ops';
+import {logicalNot} from '../ops/logical_not';
+import {mul} from '../ops/mul';
+import {zerosLike} from '../ops/tensor_ops';
+import {Tensor} from '../tensor';
+
+export const selectV2PoolGradConfig: GradConfig = {
+  kernelName: SelectV2,
+  inputsToSave: ['condition'],
+  gradFunc: (dy: Tensor, saved: Tensor[]) => {
+    const [condition] = saved;
+    return {
+      // TODO(julianoks): Return null for condition gradient
+      // when backprop supports it.
+      condition: () => cast(zerosLike(condition), 'float32'),
+      t: () => mul(dy, cast(condition, dy.dtype)),
+      e: () => mul(dy, cast(logicalNot(condition), dy.dtype))
+    };
+  }
+};

tfjs-core/src/gradients/Selu_grad.ts

@@ -18,11 +18,11 @@ import {Selu} from '../kernel_names';
 import {GradConfig} from '../kernel_registry';
 import {cast} from '../ops/array_ops';
 import {greater} from '../ops/greater';
-import {where} from '../ops/logical_ops';
 import {mul} from '../ops/mul';
 import {SELU_SCALE, SELU_SCALEALPHA} from '../ops/selu_util';
 import {scalar} from '../ops/tensor_ops';
 import {exp} from '../ops/unary_ops';
+import {where} from '../ops/where';
 import {Tensor} from '../tensor';
 
 export const seluGradConfig: GradConfig = {

tfjs-core/src/kernel_names.ts

@@ -242,6 +242,15 @@ export type LessInputs = BinaryInputs;
 export const LessEqual = 'LessEqual';
 export type LessEqualInputs = BinaryInputs;
 
+export const LogicalAnd = 'LogicalAnd';
+export type LogicalAndInputs = BinaryInputs;
+
+export const LogicalNot = 'LogicalNot';
+export type LogicalNotInputs = Pick<NamedTensorInfoMap, 'x'>;
+
+export const LogicalOr = 'LogicalOr';
+export type LogicalOrInputs = BinaryInputs;
+
 export const LRN = 'LRN';
 export type LRNInputs = Pick<NamedTensorInfoMap, 'x'>;
 export interface LRNAttrs {

tfjs-core/src/ops/band_part.ts

@@ -23,10 +23,11 @@ import {assert} from '../util';
 import {reshape, stack, unstack} from './array_ops';
 import {greaterEqual} from './greater_equal';
 import {lessEqual} from './less_equal';
-import {logicalAnd, where} from './logical_ops';
+import {logicalAnd} from './logical_and';
 import {op} from './operation';
 import {sub} from './sub';
 import {range, scalar, zeros} from './tensor_ops';
+import {where} from './where';
 
 /**
  * Copy a tensor setting everything outside a central band in each innermost

tfjs-core/src/ops/boolean_mask.ts

@@ -20,8 +20,8 @@ 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';
+import {whereAsync} from './where_async';
 
 /**
  * Apply boolean mask to tensor.

tfjs-core/src/ops/div_no_nan.ts

@@ -21,9 +21,9 @@ import {convertToTensor} from '../tensor_util_env';
 import {TensorLike} from '../types';
 
 import {div} from './div';
-import {where} from './logical_ops';
 import {op} from './operation';
 import {zerosLike} from './tensor_ops';
+import {where} from './where';
 
 /**
  * Divides two `tf.Tensor`s element-wise, A / B. Supports broadcasting. Return 0

tfjs-core/src/ops/logical_and.ts

@@ -0,0 +1,54 @@
+/**
+ * @license
+ * Copyright 2020 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 {ENGINE} from '../engine';
+import {LogicalAnd, LogicalAndInputs} from '../kernel_names';
+import {Tensor} from '../tensor';
+import {NamedTensorMap} from '../tensor_types';
+import {convertToTensor} from '../tensor_util_env';
+import {TensorLike} from '../types';
+
+import {assertAndGetBroadcastShape} from './broadcast_util';
+import {op} from './operation';
+
+/**
+ * Returns the truth value of `a AND b` element-wise. Supports broadcasting.
+ *
+ * ```js
+ * const a = tf.tensor1d([false, false, true, true], 'bool');
+ * const b = tf.tensor1d([false, true, false, true], 'bool');
+ *
+ * a.logicalAnd(b).print();
+ * ```
+ *
+ * @param a The first input tensor. Must be of dtype bool.
+ * @param b The second input tensor. Must be of dtype bool.
+ */
+/** @doc {heading: 'Operations', subheading: 'Logical'} */
+function logicalAnd_<T extends Tensor>(
+    a: Tensor|TensorLike, b: Tensor|TensorLike): T {
+  const $a = convertToTensor(a, 'a', 'logicalAnd', 'bool');
+  const $b = convertToTensor(b, 'b', 'logicalAnd', 'bool');
+  assertAndGetBroadcastShape($a.shape, $b.shape);
+
+  const inputs: LogicalAndInputs = {a: $a, b: $b};
+
+  return ENGINE.runKernelFunc(
+             backend => backend.logicalAnd($a, $b),
+             inputs as {} as NamedTensorMap, null /* grad */, LogicalAnd) as T;
+}
+
+export const logicalAnd = op({logicalAnd_});

tfjs-core/src/ops/logical_and_test.ts

@@ -0,0 +1,126 @@
+/**
+ * @license
+ * Copyright 2020 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 {expectArraysClose} from '../test_util';
+
+describeWithFlags('logicalAnd', ALL_ENVS, () => {
+  it('Tensor1D.', async () => {
+    let a = tf.tensor1d([1, 0, 0], 'bool');
+    let b = tf.tensor1d([0, 1, 0], 'bool');
+    expectArraysClose(await tf.logicalAnd(a, b).data(), [0, 0, 0]);
+
+    a = tf.tensor1d([0, 0, 0], 'bool');
+    b = tf.tensor1d([0, 0, 0], 'bool');
+    expectArraysClose(await tf.logicalAnd(a, b).data(), [0, 0, 0]);
+
+    a = tf.tensor1d([1, 1], 'bool');
+    b = tf.tensor1d([1, 1], 'bool');
+    expectArraysClose(await tf.logicalAnd(a, b).data(), [1, 1]);
+  });
+  it('mismatched Tensor1D shapes', () => {
+    const a = tf.tensor1d([1, 0], 'bool');
+    const b = tf.tensor1d([0, 1, 0], 'bool');
+    const f = () => {
+      tf.logicalAnd(a, b);
+    };
+    expect(f).toThrowError();
+  });
+
+  it('Tensor2D', async () => {
+    let a = tf.tensor2d([[1, 0, 1], [0, 0, 0]], [2, 3], 'bool');
+    let b = tf.tensor2d([[0, 0, 0], [0, 1, 0]], [2, 3], 'bool');
+    expectArraysClose(await tf.logicalAnd(a, b).data(), [0, 0, 0, 0, 0, 0]);
+
+    a = tf.tensor2d([[0, 0, 0], [1, 1, 1]], [2, 3], 'bool');
+    b = tf.tensor2d([[0, 0, 0], [1, 1, 1]], [2, 3], 'bool');
+    expectArraysClose(await tf.logicalAnd(a, b).data(), [0, 0, 0, 1, 1, 1]);
+  });
+  it('broadcasting Tensor2D shapes', async () => {
+    const a = tf.tensor2d([[1], [0]], [2, 1], 'bool');
+    const b = tf.tensor2d([[0, 1, 0], [0, 1, 0]], [2, 3], 'bool');
+    expectArraysClose(await tf.logicalAnd(a, b).data(), [0, 1, 0, 0, 0, 0]);
+  });
+
+  it('Tensor3D', async () => {
+    let a = tf.tensor3d([[[1], [0], [1]], [[0], [0], [1]]], [2, 3, 1], 'bool');
+    let b = tf.tensor3d([[[0], [0], [1]], [[1], [0], [0]]], [2, 3, 1], 'bool');
+    expectArraysClose(await tf.logicalAnd(a, b).data(), [0, 0, 1, 0, 0, 0]);
+
+    a = tf.tensor3d([[[0], [0], [0]], [[1], [1], [1]]], [2, 3, 1], 'bool');
+    b = tf.tensor3d([[[0], [0], [0]], [[1], [1], [1]]], [2, 3, 1], 'bool');
+    expectArraysClose(await tf.logicalAnd(a, b).data(), [0, 0, 0, 1, 1, 1]);
+  });
+  it('broadcasting Tensor3D shapes', async () => {
+    const a = tf.tensor3d(
+        [[[1, 0], [0, 0], [1, 1]], [[0, 0], [0, 1], [0, 0]]], [2, 3, 2],
+        'bool');
+    const b =
+        tf.tensor3d([[[0], [0], [1]], [[1], [0], [0]]], [2, 3, 1], 'bool');
+    expectArraysClose(
+        await tf.logicalAnd(a, b).data(), [0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0]);
+  });
+
+  it('Tensor4D', async () => {
+    let a = tf.tensor4d([1, 0, 1, 0], [2, 2, 1, 1], 'bool');
+    let b = tf.tensor4d([0, 1, 1, 0], [2, 2, 1, 1], 'bool');
+    expectArraysClose(await tf.logicalAnd(a, b).data(), [0, 0, 1, 0]);
+
+    a = tf.tensor4d([0, 0, 0, 0], [2, 2, 1, 1], 'bool');
+    b = tf.tensor4d([0, 0, 0, 0], [2, 2, 1, 1], 'bool');
+    expectArraysClose(await tf.logicalAnd(a, b).data(), [0, 0, 0, 0]);
+
+    a = tf.tensor4d([1, 1, 1, 1], [2, 2, 1, 1], 'bool');
+    b = tf.tensor4d([1, 1, 1, 1], [2, 2, 1, 1], 'bool');
+    expectArraysClose(await tf.logicalAnd(a, b).data(), [1, 1, 1, 1]);
+  });
+  it('broadcasting Tensor4D shapes', async () => {
+    const a = tf.tensor4d([1, 0, 1, 0], [2, 2, 1, 1], 'bool');
+    const b = tf.tensor4d(
+        [[[[1, 0]], [[0, 0]]], [[[0, 0]], [[1, 1]]]], [2, 2, 1, 2], 'bool');
+    expectArraysClose(
+        await tf.logicalAnd(a, b).data(), [1, 0, 0, 0, 0, 0, 0, 0]);
+  });
+
+  it('TensorLike', async () => {
+    const a = [true, false, false];
+    const b = [false, true, false];
+    expectArraysClose(await tf.logicalAnd(a, b).data(), [0, 0, 0]);
+  });
+
+  it('TensorLike Chained', async () => {
+    const a = tf.tensor1d([1, 0, 0], 'bool');
+    const b = [false, true, false];
+    expectArraysClose(await a.logicalAnd(b).data(), [0, 0, 0]);
+  });
+
+  it('throws when passed a as a non-tensor', () => {
+    expect(() => tf.logicalAnd({} as tf.Tensor, tf.scalar(1, 'bool')))
+        .toThrowError(/Argument 'a' passed to 'logicalAnd' must be a Tensor/);
+  });
+  it('throws when passed b as a non-tensor', () => {
+    expect(() => tf.logicalAnd(tf.scalar(1, 'bool'), {} as tf.Tensor))
+        .toThrowError(/Argument 'b' passed to 'logicalAnd' must be a Tensor/);
+  });
+
+  it('accepts a tensor-like object', async () => {
+    const a = [1, 0, 0, 1];
+    const b = [0, 1, 0, 1];
+    expectArraysClose(await tf.logicalAnd(a, b).data(), [0, 0, 0, 1]);
+  });
+});

tfjs-core/src/ops/logical_not.ts

@@ -0,0 +1,46 @@
+/**
+ * @license
+ * Copyright 2020 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 {ENGINE} from '../engine';
+import {LogicalNot, LogicalNotInputs} from '../kernel_names';
+import {Tensor} from '../tensor';
+import {NamedTensorMap} from '../tensor_types';
+import {convertToTensor} from '../tensor_util_env';
+import {TensorLike} from '../types';
+import {op} from './operation';
+
+/**
+ * Returns the truth value of `NOT x` element-wise.
+ *
+ * ```js
+ * const a = tf.tensor1d([false, true], 'bool');
+ *
+ * a.logicalNot().print();
+ * ```
+ *
+ * @param x The input tensor. Must be of dtype 'bool'.
+ */
+/** @doc {heading: 'Operations', subheading: 'Logical'} */
+function logicalNot_<T extends Tensor>(x: T|TensorLike): T {
+  const $x = convertToTensor(x, 'x', 'logicalNot', 'bool');
+  const inputs: LogicalNotInputs = {x: $x};
+  return ENGINE.runKernelFunc(
+      backend => backend.logicalNot($x), inputs as {} as NamedTensorMap,
+      null /* grad */, LogicalNot);
+}
+
+export const logicalNot = op({logicalNot_});