[WebGPU] Ensure tf.any and tf.all return bool tensors (#7928)

Fix a bug where tf.all and tf.any return int tensors instead of bool tensors.

tfjs-backend-webgpu/src/kernel_utils/reduce.ts

@@ -15,7 +15,7 @@
  * =============================================================================
  */
 
-import {backend_util, sumOutType, TensorInfo, TypedArray, util} from '@tensorflow/tfjs-core';
+import {backend_util, DataType, sumOutType, TensorInfo, TypedArray, util} from '@tensorflow/tfjs-core';
 
 import {WebGPUBackend} from '../backend_webgpu';
 import {reshape} from '../kernels/Reshape';
@@ -26,6 +26,11 @@ import {maxImplCPU} from './shared';
 import {prodImplCPU} from './shared';
 
 type ReduceTypes = 'all'|'any'|'max'|'mean'|'min'|'prod'|'sum';
+const RETURN_TYPES: {[key in ReduceTypes]?: DataType} = {
+  'mean': 'float32',
+  'all': 'bool',
+  'any': 'bool',
+};
 
 export function reduce(
     x: TensorInfo, axis: number|number[], keepDims: boolean,
@@ -79,7 +84,7 @@ export function reduce(
     const batchSize = xSize / inSize;
 
     const reduceInfo = {windowSize: inSize, inSize, batchSize, outSize: 1};
-    const dtype = reduceType === 'mean' ? 'float32' : sumOutType(x.dtype);
+    const dtype = RETURN_TYPES[reduceType] || sumOutType(x.dtype);
     const uniformData = [
       {type: 'int32', data: [inSize]},
     ];

tfjs-core/src/ops/all_test.ts

@@ -107,4 +107,10 @@ describeWithFlags('all', ALL_ENVS, () => {
         .toThrowError(
             /Argument 'x' passed to 'all' must be bool tensor, but got string/);
   });
+
+  it('returns a boolean tensor', () => {
+    const a = tf.tensor1d([1, 0], 'bool');
+    const r = tf.all(a);
+    expect(r.dtype).toEqual('bool');
+  });
 });

tfjs-core/src/ops/any_test.ts

@@ -106,4 +106,10 @@ describeWithFlags('any', ALL_ENVS, () => {
     expect(() => tf.any(['a']))
         .toThrowError(/Argument 'x' passed to 'any' must be bool tensor/);
   });
+
+  it('returns a boolean tensor', () => {
+    const a = tf.tensor1d([1, 0], 'bool');
+    const r = tf.any(a);
+    expect(r.dtype).toEqual('bool');
+  });
 });

tfjs/karma.conf.js

@@ -69,14 +69,14 @@ module.exports = function(config) {
         browser: 'chrome',
         browser_version: 'latest',
         os: 'OS X',
-        os_version: 'Sierra'
+        os_version: 'High Sierra'
       },
       bs_firefox_mac: {
         base: 'BrowserStack',
         browser: 'firefox',
-        browser_version: 'latest',
+        browser_version: '90',
         os: 'OS X',
-        os_version: 'Sierra'
+        os_version: 'High Sierra'
       },
       chrome_with_swift_shader: {
         base: 'Chrome',