[wasm] Add split, sqrt kernels.

tfjs-backend-wasm/src/cc/kernels/Sqrt.cc

@@ -0,0 +1,36 @@
+/* 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.
+ * ===========================================================================*/
+
+#ifdef __EMSCRIPTEN__
+#include <emscripten.h>
+#endif
+
+#include <math.h>
+
+#include "src/cc/backend.h"
+#include "src/cc/unary.h"
+
+namespace tfjs {
+namespace wasm {
+// We use C-style API to interface with Javascript.
+extern "C" {
+
+#ifdef __EMSCRIPTEN__
+EMSCRIPTEN_KEEPALIVE
+#endif
+void Sqrt(const int x_id, const int out_id) { unary(x_id, out_id, sqrt); }
+
+}  // extern "C"
+}  // namespace wasm
+}  // namespace tfjs

tfjs-backend-wasm/src/kernels/Split.ts

@@ -0,0 +1,55 @@
+/**
+ * @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 {NamedAttrMap, NamedTensorInfoMap, registerKernel, SplitV, SplitVAttrs, SplitVInputs, util} from '@tensorflow/tfjs-core';
+
+import {BackendWasm} from '../backend_wasm';
+
+import {slice} from './Slice';
+
+export function split(args: {
+  inputs: NamedTensorInfoMap,
+  attrs: NamedAttrMap,
+  backend: BackendWasm
+}) {
+  const {inputs, attrs, backend} = args;
+  const {x} = inputs as {} as SplitVInputs;
+  const {numOrSizeSplits, axis} = attrs as {} as SplitVAttrs;
+
+  const $axis = util.parseAxisParam(axis, x.shape)[0];
+
+  let splitSizes: number[];
+  if (typeof (numOrSizeSplits) === 'number') {
+    splitSizes =
+        new Array(numOrSizeSplits).fill(x.shape[$axis] / numOrSizeSplits);
+  } else {
+    splitSizes = numOrSizeSplits;
+  }
+
+  const begin = new Array(x.shape.length).fill(0);
+  const size = x.shape.slice();
+  return splitSizes.map(s => {
+    const xSliceSize = [...size];
+    xSliceSize[$axis] = s;
+    const xSlice =
+        slice({inputs: {x}, attrs: {begin, size: xSliceSize}, backend});
+    begin[$axis] += s;
+    return xSlice;
+  });
+}
+
+registerKernel({kernelName: SplitV, backendName: 'wasm', kernelFunc: split});

tfjs-backend-wasm/src/kernels/Sqrt.ts

@@ -0,0 +1,19 @@
+/**
+ * @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 {registerUnaryKernel} from './unary_kernel';
+registerUnaryKernel('Sqrt');

tfjs-backend-wasm/src/kernels/all_kernels.ts

@@ -70,6 +70,8 @@ import './Sigmoid';
 import './Sin';
 import './Slice';
 import './Softmax';
+import './Split';
+import './Sqrt';
 import './Square';
 import './Sub';
 import './Sum';

tfjs-backend-wasm/src/setup_test.ts

@@ -224,6 +224,7 @@ const TEST_FILTERS: TestFilter[] = [
     include: 'transpose',
     excludes: ['oneHot']  // oneHot not yet implemented.
   },
+  {include: 'split'},
   {include: 'pad ', excludes: ['complex', 'zerosLike']},
   {include: 'clip', excludes: ['gradient']},
   {include: 'addN'},
@@ -328,6 +329,10 @@ const TEST_FILTERS: TestFilter[] = [
     startsWith: 'rsqrt ',
     excludes: ['gradient']  // Gradient not yet implemented.
   },
+  {
+    startsWith: 'sqrt ',
+    excludes: ['gradient']  // Gradient not yet implemented.
+  },
   {
     startsWith: 'zerosLike',
     // Complex numbers not supported yet.

tfjs-core/src/backends/split_shared.ts

@@ -17,14 +17,18 @@
 
 import {Tensor} from '../tensor';
 
+// TODO(annxingyuan): Use this helper in WASM Split kernel once intermediate
+// kernels have been modularized in WebGL and CPU
+// https://github.com/tensorflow/tfjs/issues/2822.
 /** Shared implementation of the split kernel across WebGL and CPU. */
 export function split<T extends Tensor>(
     x: T, sizeSplits: number[], axis: number): T[] {
   const begin = new Array(x.rank).fill(0);
   const size = x.shape.slice();
   return sizeSplits.map(s => {
-    size[axis] = s;
-    const slice = x.slice(begin, size);
+    const sliceSize = [...size];
+    sliceSize[axis] = s;
+    const slice = x.slice(begin, sliceSize);
     begin[axis] += s;
     return slice;
   });

tfjs-core/src/ops/split.ts

@@ -80,7 +80,6 @@ function split_<T extends Tensor>(
   }
 
   const forward: ForwardFunc<Tensor> = (backend, _) => {
-    const $axis = parseAxisParam(axis, $x.shape)[0];
     return backend.split($x, splitSizes, $axis) as {} as T;
   };
 

tfjs-core/src/ops/unary_ops.ts

@@ -325,13 +325,13 @@ function sqrt_<T extends Tensor>(x: T|TensorLike): T {
 
   const grad = (dy: T, saved: Tensor[]) => {
     const [$x] = saved;
-    return {$x: () => dy.div($x.toFloat().sqrt().mul(2))} as {$x: () => T};
+    return {x: () => dy.div($x.toFloat().sqrt().mul(2))} as {x: () => T};
   };
   return ENGINE.runKernelFunc((backend, save) => {
     const res = backend.sqrt($x);
     save([$x]);
     return res;
-  }, {$x}, grad);
+  }, {x: $x}, grad, 'Sqrt', {});
 }
 
 /**