[wasm] Use export KernelConfig pattern in WASM.

tfjs-backend-wasm/src/index.ts

@@ -15,7 +15,6 @@
  * =============================================================================
  */
 
-import './kernels/all_kernels';
 import './register_all_kernels';
 
 export {BackendWasm, setWasmPath} from './backend_wasm';

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

@@ -14,6 +14,7 @@
  * limitations under the License.
  * =============================================================================
  */
+import {KernelConfig} from '@tensorflow/tfjs-core';
 
-import {registerUnaryKernel} from './unary_kernel';
-registerUnaryKernel('Abs');
+import {createUnaryKernelConfig} from './unary_kernel';
+export const absConfig: KernelConfig = createUnaryKernelConfig('Abs');

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

@@ -15,6 +15,11 @@
  * =============================================================================
  */
 
-import {registerBinaryKernel} from './binary_kernel';
+import {KernelConfig} from '@tensorflow/tfjs-core';
+
+import {createBinaryKernelConfig} from './binary_kernel';
+
 const supportsFullBroadcast = true;
-registerBinaryKernel('Add', supportsFullBroadcast);
+
+export const addConfig: KernelConfig =
+    createBinaryKernelConfig('Add', supportsFullBroadcast);

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

@@ -15,9 +15,10 @@
  * =============================================================================
  */
 
-import {KernelFunc, registerKernel, TensorInfo, util} from '@tensorflow/tfjs-core';
+import {KernelConfig, KernelFunc, TensorInfo, util} from '@tensorflow/tfjs-core';
 
 import {BackendWasm} from '../backend_wasm';
+
 import {CppDType} from './types';
 
 let wasmFunc:
@@ -50,9 +51,9 @@ function addn(args: {inputs: TensorInfo[], backend: BackendWasm}) {
   return out;
 }
 
-registerKernel({
+export const addNConfig: KernelConfig = {
   kernelName: 'AddN',
   backendName: 'wasm',
   setupFunc,
   kernelFunc: addn as {} as KernelFunc,
-});
+};

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

@@ -15,7 +15,7 @@
  * =============================================================================
  */
 
-import {ArgMax, ArgMaxAttrs, ArgMaxInputs, KernelFunc, registerKernel, util} from '@tensorflow/tfjs-core';
+import {ArgMax, ArgMaxAttrs, ArgMaxInputs, KernelConfig, KernelFunc, util} from '@tensorflow/tfjs-core';
 
 import {BackendWasm} from '../backend_wasm';
 
@@ -73,9 +73,9 @@ function argmax(
   return out;
 }
 
-registerKernel({
+export const argMaxConfig: KernelConfig = {
   kernelName: ArgMax,
   backendName: 'wasm',
   kernelFunc: argmax as {} as KernelFunc,
   setupFunc: setup
-});
+};

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

@@ -15,7 +15,7 @@
  * =============================================================================
  */
 
-import {AvgPoolAttrs, AvgPoolInputs, backend_util, KernelFunc, registerKernel, Tensor4D} from '@tensorflow/tfjs-core';
+import {AvgPoolAttrs, AvgPoolInputs, backend_util, KernelConfig, KernelFunc, Tensor4D} from '@tensorflow/tfjs-core';
 
 import {BackendWasm} from '../backend_wasm';
 
@@ -87,9 +87,9 @@ function avgPool(
   return out;
 }
 
-registerKernel({
+export const avgPoolConfig: KernelConfig = {
   kernelName: 'AvgPool',
   backendName: 'wasm',
   setupFunc: setup,
   kernelFunc: avgPool as {} as KernelFunc
-});
+};

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

@@ -15,7 +15,7 @@
  * =============================================================================
  */
 
-import {NamedAttrMap, NamedTensorInfoMap, registerKernel, TensorInfo} from '@tensorflow/tfjs-core';
+import {KernelConfig, NamedAttrMap, NamedTensorInfoMap, TensorInfo} from '@tensorflow/tfjs-core';
 
 import {BackendWasm} from '../backend_wasm';
 
@@ -82,9 +82,9 @@ function batchMatMul(args: {
   return out;
 }
 
-registerKernel({
+export const batchMatMulConfig: KernelConfig = {
   kernelName: 'BatchMatMul',
   backendName: 'wasm',
   setupFunc: setup,
   kernelFunc: batchMatMul
-});
+};

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

@@ -15,7 +15,7 @@
  * =============================================================================
  */
 
-import {DataType, NamedAttrMap, NamedTensorInfoMap, registerKernel} from '@tensorflow/tfjs-core';
+import {DataType, KernelConfig, NamedAttrMap, NamedTensorInfoMap} from '@tensorflow/tfjs-core';
 import {TensorInfo} from '@tensorflow/tfjs-core';
 
 import {BackendWasm} from '../backend_wasm';
@@ -39,8 +39,8 @@ export function cast(
   return out;
 }
 
-registerKernel({
+export const castConfig: KernelConfig = {
   kernelName: 'Cast',
   backendName: 'wasm',
   kernelFunc: cast,
-});
+};

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

@@ -15,7 +15,7 @@
  * =============================================================================
  */
 
-import {NamedAttrMap, NamedTensorInfoMap, registerKernel} from '@tensorflow/tfjs-core';
+import {KernelConfig, NamedAttrMap, NamedTensorInfoMap} from '@tensorflow/tfjs-core';
 import {TensorInfo} from '@tensorflow/tfjs-core';
 
 import {BackendWasm} from '../backend_wasm';
@@ -55,9 +55,9 @@ function clip(args: {
   return out;
 }
 
-registerKernel({
+export const clipByValueConfig: KernelConfig = {
   kernelName: 'ClipByValue',
   backendName: 'wasm',
   setupFunc: setup,
   kernelFunc: clip
-});
+};

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

@@ -15,7 +15,7 @@
  * =============================================================================
  */
 
-import {backend_util, KernelFunc, NamedAttrMap, registerKernel, TensorInfo, util} from '@tensorflow/tfjs-core';
+import {backend_util, KernelConfig, KernelFunc, NamedAttrMap, TensorInfo, util} from '@tensorflow/tfjs-core';
 
 import {BackendWasm} from '../backend_wasm';
 
@@ -54,8 +54,8 @@ function concat(
   return out;
 }
 
-registerKernel({
+export const concatConfig: KernelConfig = {
   kernelName: 'Concat',
   backendName: 'wasm',
   kernelFunc: concat as {} as KernelFunc,
-});
+};

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

@@ -15,7 +15,7 @@
  * =============================================================================
  */
 
-import {backend_util, Conv2DAttrs, KernelFunc, NamedTensorInfoMap, registerKernel, Tensor4D, TensorInfo} from '@tensorflow/tfjs-core';
+import {backend_util, Conv2DAttrs, KernelConfig, KernelFunc, NamedTensorInfoMap, Tensor4D, TensorInfo} from '@tensorflow/tfjs-core';
 
 import {BackendWasm} from '../backend_wasm';
 
@@ -100,9 +100,9 @@ function conv2d(
   return out;
 }
 
-registerKernel({
+export const conv2DConfig: KernelConfig = {
   kernelName: 'Conv2D',
   backendName: 'wasm',
   setupFunc: setup,
   kernelFunc: conv2d as {} as KernelFunc
-});
+};

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

@@ -15,7 +15,7 @@
  * =============================================================================
  */
 
-import {backend_util, Conv2DBackpropInput, Conv2DBackpropInputAttrs, Conv2DBackpropInputInputs, NamedAttrMap, NamedTensorInfoMap, registerKernel, TensorInfo, util} from '@tensorflow/tfjs-core';
+import {backend_util, Conv2DBackpropInput, Conv2DBackpropInputAttrs, Conv2DBackpropInputInputs, KernelConfig, NamedAttrMap, NamedTensorInfoMap, TensorInfo, util} from '@tensorflow/tfjs-core';
 
 import {BackendWasm} from '../backend_wasm';
 
@@ -121,9 +121,9 @@ function conv2DBackpropInput(args: {
   return out;
 }
 
-registerKernel({
+export const conv2DBackpropInputConfig: KernelConfig = {
   kernelName: Conv2DBackpropInput,
   backendName: 'wasm',
   setupFunc: setup,
   kernelFunc: conv2DBackpropInput
-});
+};

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

@@ -15,5 +15,8 @@
  * =============================================================================
  */
 
-import {registerUnaryKernel} from './unary_kernel';
-registerUnaryKernel('Cos');
+import {KernelConfig} from '@tensorflow/tfjs-core';
+
+import {createUnaryKernelConfig} from './unary_kernel';
+
+export const cosConfig: KernelConfig = createUnaryKernelConfig('Cos');

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

@@ -15,7 +15,7 @@
  * =============================================================================
  */
 
-import {CropAndResize, CropAndResizeAttrs, CropAndResizeInputs, NamedAttrMap, NamedTensorInfoMap, registerKernel, TensorInfo} from '@tensorflow/tfjs-core';
+import {CropAndResize, CropAndResizeAttrs, CropAndResizeInputs, KernelConfig, NamedAttrMap, NamedTensorInfoMap, TensorInfo} from '@tensorflow/tfjs-core';
 
 import {BackendWasm} from '../backend_wasm';
 
@@ -91,9 +91,9 @@ function cropAndResize(args: {
   return out;
 }
 
-registerKernel({
+export const cropAndResizeConfig: KernelConfig = {
   kernelName: CropAndResize,
   backendName: 'wasm',
   setupFunc: setup,
   kernelFunc: cropAndResize
-});
+};

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

@@ -15,7 +15,7 @@
  * =============================================================================
  */
 
-import {backend_util, DepthwiseConv2dNativeAttrs, KernelFunc, NamedTensorInfoMap, registerKernel, Tensor4D, TensorInfo} from '@tensorflow/tfjs-core';
+import {backend_util, DepthwiseConv2dNativeAttrs, KernelConfig, KernelFunc, NamedTensorInfoMap, Tensor4D, TensorInfo} from '@tensorflow/tfjs-core';
 
 import {BackendWasm} from '../backend_wasm';
 
@@ -107,9 +107,9 @@ function depthwiseConv2d(args: {
   return out;
 }
 
-registerKernel({
+export const depthwiseConv2DNativeConfig: KernelConfig = {
   kernelName: 'DepthwiseConv2dNative',
   backendName: 'wasm',
   setupFunc: setup,
   kernelFunc: depthwiseConv2d as {} as KernelFunc
-});
+};

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

@@ -15,6 +15,9 @@
  * =============================================================================
  */
 
-import {registerBinaryKernel} from './binary_kernel';
+import {KernelConfig} from '@tensorflow/tfjs-core';
+
+import {createBinaryKernelConfig} from './binary_kernel';
 const supportsFullBroadcast = true;
-registerBinaryKernel('Div', supportsFullBroadcast);
+export const divConfig: KernelConfig =
+    createBinaryKernelConfig('Div', supportsFullBroadcast);

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

@@ -14,7 +14,9 @@
  * limitations under the License.
  * =============================================================================
  */
+import {KernelConfig} from '@tensorflow/tfjs-core';
 
-import {registerBinaryKernel} from './binary_kernel';
+import {createBinaryKernelConfig} from './binary_kernel';
 const supportsFullBroadcast = false;
-registerBinaryKernel('Equal', supportsFullBroadcast, 'bool');
+export const equalConfig: KernelConfig =
+    createBinaryKernelConfig('Equal', supportsFullBroadcast, 'bool');

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

@@ -14,6 +14,7 @@
  * limitations under the License.
  * =============================================================================
  */
+import {KernelConfig} from '@tensorflow/tfjs-core';
 
-import {registerUnaryKernel} from './unary_kernel';
-registerUnaryKernel('Exp');
+import {createUnaryKernelConfig} from './unary_kernel';
+export const expConfig: KernelConfig = createUnaryKernelConfig('Exp');

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

@@ -15,7 +15,7 @@
  * =============================================================================
  */
 
-import {KernelFunc, registerKernel} from '@tensorflow/tfjs-core';
+import {KernelConfig, KernelFunc} from '@tensorflow/tfjs-core';
 import {Fill, FillAttrs} from '@tensorflow/tfjs-core';
 
 import {BackendWasm} from '../backend_wasm';
@@ -28,8 +28,8 @@ function fill(args: {attrs: FillAttrs, backend: BackendWasm}) {
   return out;
 }
 
-registerKernel({
+export const fillConfig: KernelConfig = {
   kernelName: Fill,
   backendName: 'wasm',
   kernelFunc: fill as {} as KernelFunc,
-});
+};

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

@@ -15,6 +15,9 @@
  * =============================================================================
  */
 
-import {registerBinaryKernel} from './binary_kernel';
+import {KernelConfig} from '@tensorflow/tfjs-core';
+
+import {createBinaryKernelConfig} from './binary_kernel';
 const supportsFullBroadcast = false;
-registerBinaryKernel('FloorDiv', supportsFullBroadcast);
+export const floorDivConfig: KernelConfig =
+    createBinaryKernelConfig('FloorDiv', supportsFullBroadcast);

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

@@ -15,7 +15,7 @@
  * =============================================================================
  */
 
-import {NamedAttrMap, NamedTensorInfoMap, registerKernel, TensorInfo, util} from '@tensorflow/tfjs-core';
+import {KernelConfig, NamedAttrMap, NamedTensorInfoMap, TensorInfo, util} from '@tensorflow/tfjs-core';
 
 import {BackendWasm} from '../backend_wasm';
 
@@ -67,9 +67,9 @@ function fusedBatchNorm(
   return out;
 }
 
-registerKernel({
+export const fusedBatchNormConfig: KernelConfig = {
   kernelName: 'FusedBatchNorm',
   backendName: 'wasm',
   setupFunc: setup,
   kernelFunc: fusedBatchNorm
-});
+};

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

@@ -15,7 +15,7 @@
  * =============================================================================
  */
 
-import {backend_util, KernelFunc, NamedTensorInfoMap, registerKernel, TensorInfo} from '@tensorflow/tfjs-core';
+import {backend_util, KernelConfig, KernelFunc, NamedTensorInfoMap, TensorInfo} from '@tensorflow/tfjs-core';
 
 import {BackendWasm} from '../backend_wasm';
 
@@ -136,9 +136,9 @@ function fusedConv2d(args: {
   return out;
 }
 
-registerKernel({
+export const fusedConv2DConfig: KernelConfig = {
   kernelName: 'FusedConv2D',
   backendName: 'wasm',
   setupFunc: setup,
   kernelFunc: fusedConv2d as {} as KernelFunc
-});
+};