Add Cos to kernel registry (cpu backend) (#3907)

Author: jinjingforever

tfjs-backend-cpu/.vscode/launch.json

@@ -0,0 +1,18 @@
+{
+  // Use IntelliSense to learn about possible attributes.
+  // Hover to view descriptions of existing attributes.
+  // For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387
+  "version": "0.2.0",
+  "configurations": [
+    {
+      "name": "Run tests",
+      "type": "node",
+      "request": "launch",
+      "runtimeExecutable": "yarn",
+      "runtimeArgs": ["test"],
+      "sourceMaps": true,
+      "cwd": "${workspaceRoot}",
+      "protocol": "inspector",
+    }
+  ]
+}

tfjs-backend-cpu/src/backend_cpu.ts

@@ -1279,17 +1279,6 @@ export class MathBackendCPU extends KernelBackend {
     return this.makeOutput(resultValues, x.shape, 'float32');
   }
 
-  cos<T extends Tensor>(x: T): T {
-    assertNotComplex(x, 'cos');
-
-    const resultValues = new Float32Array(x.size);
-    const values = this.readSync(x.dataId) as TypedArray;
-    for (let i = 0; i < values.length; ++i) {
-      resultValues[i] = Math.cos(values[i]);
-    }
-    return this.makeOutput(resultValues, x.shape, 'float32');
-  }
-
   tan<T extends Tensor>(x: T): T {
     assertNotComplex(x, 'tan');
 

tfjs-backend-cpu/src/kernels/Cos.ts

@@ -0,0 +1,40 @@
+/**
+ * @license
+ * Copyright 2020 Google LLC. 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 {Cos, CosInputs, KernelConfig, TypedArray, util} from '@tensorflow/tfjs-core';
+
+import {MathBackendCPU} from '../backend_cpu';
+import {assertNotComplex} from '../cpu_util';
+
+export const cosConfig: KernelConfig = {
+  kernelName: Cos,
+  backendName: 'cpu',
+  kernelFunc: ({inputs, backend}) => {
+    const {x} = inputs as CosInputs;
+    const cpuBackend = backend as MathBackendCPU;
+    assertNotComplex(x, 'cos');
+
+    const values = cpuBackend.data.get(x.dataId).values as TypedArray;
+    const xSize = util.sizeFromShape(x.shape);
+    const newValues = new Float32Array(xSize);
+    for (let i = 0; i < xSize; ++i) {
+      newValues[i] = Math.cos(values[i]);
+    }
+    const dataId = cpuBackend.write(newValues, x.shape, x.dtype);
+    return {dataId, shape: x.shape, dtype: x.dtype};
+  }
+};

tfjs-backend-cpu/src/register_all_kernels.ts

@@ -19,6 +19,7 @@
 // the contents of this file and import only the kernels that are needed.
 import {KernelConfig, registerKernel} from '@tensorflow/tfjs-core';
 
+import {cosConfig} from './kernels/Cos';
 import {dilation2dConfig} from './kernels/Dilation2D';
 import {dilation2dBackpropFilterConfig} from './kernels/Dilation2DBackpropFilter';
 import {dilation2dBackpropInputConfig} from './kernels/Dilation2DBackpropInput';
@@ -39,7 +40,7 @@ import {transposeConfig} from './kernels/Transpose';
 
 // List all kernel configs here
 const kernelConfigs: KernelConfig[] = [
-  dilation2dConfig, dilation2dBackpropInputConfig,
+  cosConfig, dilation2dConfig, dilation2dBackpropInputConfig,
   dilation2dBackpropFilterConfig, divConfig, flipLeftRightConfig,
   identityConfig, maxPoolWithArgmaxConfig, maxConfig, nonMaxSuppressionV4Config,
   nonMaxSuppressionV5Config, padV2Config, reshapeConfig, rotateWithOffsetConfig,