From 7af28b17402393ac5074974cd5370b9ab5d73520 Mon Sep 17 00:00:00 2001
From: Na Li <linazhao@google.com>
Date: Thu, 27 Aug 2020 12:16:02 -0700
Subject: [PATCH] Modularize kernel identity.

---
 tfjs-backend-cpu/src/kernels/Identity.ts     | 36 ++++++++++++++++++++
 tfjs-backend-cpu/src/kernels/Reshape.ts      |  2 +-
 tfjs-backend-cpu/src/register_all_kernels.ts |  3 +-
 3 files changed, 39 insertions(+), 2 deletions(-)
 create mode 100644 tfjs-backend-cpu/src/kernels/Identity.ts

diff --git a/tfjs-backend-cpu/src/kernels/Identity.ts b/tfjs-backend-cpu/src/kernels/Identity.ts
new file mode 100644
index 00000000000..a5cff6c996a
--- /dev/null
+++ b/tfjs-backend-cpu/src/kernels/Identity.ts
@@ -0,0 +1,36 @@
+/**
+ * @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 {Identity, IdentityInputs, KernelConfig, KernelFunc, TensorInfo} from '@tensorflow/tfjs-core';
+
+import {MathBackendCPU} from '../backend_cpu';
+
+export function identity(
+    args: {inputs: IdentityInputs, backend: MathBackendCPU}): TensorInfo {
+  const {inputs, backend} = args;
+  const {x} = inputs;
+
+  backend.incRef(x.dataId);
+
+  return {dataId: x.dataId, shape: x.shape, dtype: x.dtype};
+}
+
+export const identityConfig: KernelConfig = {
+  kernelName: Identity,
+  backendName: 'cpu',
+  kernelFunc: identity as {} as KernelFunc
+};
diff --git a/tfjs-backend-cpu/src/kernels/Reshape.ts b/tfjs-backend-cpu/src/kernels/Reshape.ts
index e7d5e359266..e9a878b6346 100644
--- a/tfjs-backend-cpu/src/kernels/Reshape.ts
+++ b/tfjs-backend-cpu/src/kernels/Reshape.ts
@@ -1,6 +1,6 @@
 /**
  * @license
- * Copyright 2019 Google LLC. All Rights Reserved.
+ * 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
diff --git a/tfjs-backend-cpu/src/register_all_kernels.ts b/tfjs-backend-cpu/src/register_all_kernels.ts
index 70ec59226df..734ae7034ff 100644
--- a/tfjs-backend-cpu/src/register_all_kernels.ts
+++ b/tfjs-backend-cpu/src/register_all_kernels.ts
@@ -24,6 +24,7 @@ import {dilation2dBackpropFilterConfig} from './kernels/Dilation2DBackpropFilter
 import {dilation2dBackpropInputConfig} from './kernels/Dilation2DBackpropInput';
 import {divConfig} from './kernels/Div';
 import {flipLeftRightConfig} from './kernels/FlipLeftRight';
+import {identityConfig} from './kernels/Identity';
 import {maxConfig} from './kernels/Max';
 import {maxPoolWithArgmaxConfig} from './kernels/MaxPoolWithArgmax';
 import {nonMaxSuppressionV4Config} from './kernels/NonMaxSuppressionV4';
@@ -40,7 +41,7 @@ import {transposeConfig} from './kernels/Transpose';
 const kernelConfigs: KernelConfig[] = [
   dilation2dConfig, dilation2dBackpropInputConfig,
   dilation2dBackpropFilterConfig, divConfig, flipLeftRightConfig,
-  maxPoolWithArgmaxConfig, maxConfig, nonMaxSuppressionV4Config,
+  identityConfig, maxPoolWithArgmaxConfig, maxConfig, nonMaxSuppressionV4Config,
   nonMaxSuppressionV5Config, padV2Config, reshapeConfig, rotateWithOffsetConfig,
   spaceToBatchNDConfig, squareConfig, squaredDifferenceConfig, transposeConfig
 ];