Add range kernel to wasm and node

tfjs-backend-wasm/src/kernel_utils/shared.ts

@@ -19,5 +19,7 @@
 // side effects.
 // tslint:disable-next-line: no-imports-from-dist
 import {sliceImpl as sliceImplCPU} from '@tensorflow/tfjs-backend-cpu/dist/shared';
+// tslint:disable-next-line: no-imports-from-dist
+import {rangeImpl as rangeImplCPU} from '@tensorflow/tfjs-backend-cpu/dist/shared';
 
-export {sliceImplCPU};
+export {rangeImplCPU, sliceImplCPU};

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

@@ -0,0 +1,39 @@
+/**
+ * @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 {KernelConfig, KernelFunc, Range, RangeAttrs, TensorInfo} from '@tensorflow/tfjs-core';
+
+import {BackendWasm} from '../backend_wasm';
+import {rangeImplCPU} from '../kernel_utils/shared';
+
+export const range =
+    (args: {backend: BackendWasm, attrs: RangeAttrs}): TensorInfo => {
+      const {backend, attrs} = args;
+      const {start, stop, step, dtype} = attrs;
+      const values = rangeImplCPU(start, stop, step, dtype);
+
+      const out = backend.makeOutput([values.length], dtype);
+      const outVals = backend.typedArrayFromHeap(out);
+      outVals.set(values);
+      return out;
+    };
+
+export const rangeConfig: KernelConfig = {
+  kernelName: Range,
+  backendName: 'wasm',
+  kernelFunc: range as {} as KernelFunc
+};

tfjs-backend-wasm/src/register_all_kernels.ts

@@ -75,6 +75,7 @@ import {padV2Config} from './kernels/PadV2';
 import {powConfig} from './kernels/Pow';
 import {preluConfig} from './kernels/Prelu';
 import {prodConfig} from './kernels/Prod';
+import {rangeConfig} from './kernels/Range';
 import {realDivConfig} from './kernels/RealDiv';
 import {reluConfig} from './kernels/Relu';
 import {relu6Config} from './kernels/Relu6';
@@ -161,6 +162,7 @@ const kernelConfigs: KernelConfig[] = [
   powConfig,
   preluConfig,
   prodConfig,
+  rangeConfig,
   realDivConfig,
   reluConfig,
   relu6Config,

tfjs-core/src/ops/range.ts

@@ -15,14 +15,10 @@
  * =============================================================================
  */
 
-import {ENGINE, ForwardFunc} from '../engine';
+import {ENGINE} from '../engine';
 import {Range, RangeAttrs} from '../kernel_names';
 import {NamedAttrMap} from '../kernel_registry';
-import {Tensor, Tensor1D} from '../tensor';
-import {makeZerosTypedArray} from '../util';
-
-import {tensor1d} from './tensor1d';
-import {zeros} from './zeros';
+import {Tensor1D} from '../tensor';
 
 /**
  * Creates a new `tf.Tensor1D` filled with the numbers in the range provided.
@@ -50,36 +46,7 @@ export function range(
     throw new Error('Cannot have a step of zero');
   }
 
-  const forward: ForwardFunc<Tensor> = () => {
-    const sameStartStop = start === stop;
-    const increasingRangeNegativeStep = start < stop && step < 0;
-    const decreasingRangePositiveStep = stop < start && step > 1;
-
-    if (sameStartStop || increasingRangeNegativeStep ||
-        decreasingRangePositiveStep) {
-      return zeros([0], dtype);
-    }
-
-    const numElements = Math.abs(Math.ceil((stop - start) / step));
-    const values = makeZerosTypedArray(numElements, dtype);
-
-    if (stop < start && step === 1) {
-      // Auto adjust the step's sign if it hasn't been set
-      // (or was set to 1)
-      step = -1;
-    }
-
-    values[0] = start;
-    for (let i = 1; i < values.length; i++) {
-      values[i] = values[i - 1] + step;
-    }
-
-    return tensor1d(values, dtype);
-  };
-
   const attrs: RangeAttrs = {start, stop, step, dtype};
 
-  return ENGINE.runKernelFunc(
-             forward, {} /* inputs */, null /* grad */, Range,
-             attrs as {} as NamedAttrMap) as Tensor1D;
+  return ENGINE.runKernel(Range, {} /* inputs */, attrs as {} as NamedAttrMap);
 }

tfjs-node/src/kernels/Range.ts

@@ -0,0 +1,61 @@
+/**
+ * @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 {KernelConfig, Range, RangeAttrs, scalar, zeros} from '@tensorflow/tfjs';
+
+import {createTensorsTypeOpAttr, NodeJSKernelBackend} from '../nodejs_kernel_backend';
+
+export const rangeConfig: KernelConfig = {
+  kernelName: Range,
+  backendName: 'tensorflow',
+  kernelFunc: (args) => {
+    const backend = args.backend as NodeJSKernelBackend;
+    const {start, stop, dtype} = args.attrs as {} as RangeAttrs;
+    let {step} = args.attrs as {} as RangeAttrs;
+
+    // TensorFlow.js specific allowances
+    const sameStartStop = start === stop;
+    const increasingRangeNegativeStep = start < stop && step < 0;
+    const decreasingRangePositiveStep = stop < start && step > 1;
+
+    if (sameStartStop || increasingRangeNegativeStep ||
+        decreasingRangePositiveStep) {
+      return zeros([0], dtype);
+    }
+
+    if (stop < start && step === 1) {
+      // Auto adjust the step's sign if it hasn't been set
+      // (or was set to 1)
+      step = -1;
+    }
+
+    const opAttrs = [createTensorsTypeOpAttr('Tidx', dtype)];
+    const startTensor = scalar(start);
+    const stopTensor = scalar(stop);
+    const stepTensor = scalar(step);
+    const res = backend.executeSingleOutput(
+        Range, opAttrs, [startTensor, stopTensor, stepTensor]);
+    const castedRes = res.cast(dtype);
+
+    startTensor.dispose();
+    stopTensor.dispose();
+    stepTensor.dispose();
+    res.dispose();
+
+    return castedRes;
+  }
+};

tfjs-node/src/register_all_kernels.ts

@@ -125,6 +125,7 @@ import {padV2Config} from './kernels/PadV2';
 import {powConfig} from './kernels/Pow';
 import {preluConfig} from './kernels/Prelu';
 import {prodConfig} from './kernels/Prod';
+import {rangeConfig} from './kernels/Range';
 import {realConfig} from './kernels/Real';
 import {realDivConfig} from './kernels/RealDiv';
 import {reciprocalConfig} from './kernels/Reciprocal';
@@ -275,6 +276,7 @@ const kernelConfigs: KernelConfig[] = [
   powConfig,
   preluConfig,
   prodConfig,
+  rangeConfig,
   realConfig,
   realDivConfig,
   reciprocalConfig,