Add StaticRegexReplace Op (#7379)

Author: mattsoulanille

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

@@ -0,0 +1,37 @@
+/**
+ * @license
+ * Copyright 2023 Google LLC.
+ * 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, StaticRegexReplace, StaticRegexReplaceAttrs} from '@tensorflow/tfjs-core';
+import {createSimpleUnaryImpl} from '../utils/unary_impl';
+import {unaryKernelFuncFromImpl} from '../utils/unary_utils';
+
+export const staticRegexReplaceImpl = createSimpleUnaryImpl<string,
+  string>((x: string, attrs) => {
+    const {pattern, replaceGlobal, rewrite} =
+      attrs as unknown as StaticRegexReplaceAttrs;
+    // TODO(mattSoulanille): Don't create a regex each time.
+    return x.replace(new RegExp(pattern, replaceGlobal ? 'g' : ''), rewrite);
+});
+
+const staticRegexReplace =
+  unaryKernelFuncFromImpl(StaticRegexReplace, staticRegexReplaceImpl);
+
+export const staticRegexReplaceConfig: KernelConfig = {
+  kernelName: StaticRegexReplace,
+  backendName: 'cpu',
+  kernelFunc: staticRegexReplace,
+};

tfjs-backend-cpu/src/register_all_kernels.ts

@@ -170,6 +170,7 @@ import {splitVConfig} from './kernels/SplitV';
 import {sqrtConfig} from './kernels/Sqrt';
 import {squareConfig} from './kernels/Square';
 import {squaredDifferenceConfig} from './kernels/SquaredDifference';
+import {staticRegexReplaceConfig} from './kernels/StaticRegexReplace';
 import {stepConfig} from './kernels/Step';
 import {stridedSliceConfig} from './kernels/StridedSlice';
 import {stringNGramsConfig} from './kernels/StringNGrams';
@@ -342,6 +343,7 @@ const kernelConfigs: KernelConfig[] = [
   sqrtConfig,
   squareConfig,
   squaredDifferenceConfig,
+  staticRegexReplaceConfig,
   stepConfig,
   stridedSliceConfig,
   stringNGramsConfig,

tfjs-backend-cpu/src/shared.ts

@@ -55,6 +55,7 @@ export {sparseReshapeImpl} from './kernels/SparseReshape_impl';
 export {sparseSegmentReductionImpl} from './kernels/SparseSegmentReduction_impl';
 export {sqrtImpl} from './kernels/Sqrt';
 export {squaredDifferenceImpl} from './kernels/SquaredDifference';
+export {staticRegexReplaceImpl} from './kernels/StaticRegexReplace';
 export {stridedSliceImpl} from './kernels/StridedSlice_impl';
 export {stringNGramsImpl} from './kernels/StringNGrams_impl';
 export {stringSplitImpl} from './kernels/StringSplit_impl';

tfjs-backend-cpu/src/utils/unary_impl.ts

@@ -15,18 +15,19 @@
  * =============================================================================
  */
 
-import {NumericDataType, util} from '@tensorflow/tfjs-core';
+import {util} from '@tensorflow/tfjs-core';
 
 import {SimpleUnaryImpl, SimpleUnaryOperation} from './unary_types';
 
 /**
  * Template that creates implementation for unary op.
  */
-export function createSimpleUnaryImpl(op: SimpleUnaryOperation):
-    SimpleUnaryImpl {
+export function createSimpleUnaryImpl<I extends number | string = number,
+  O extends number | string = number>(op: SimpleUnaryOperation<I, O>):
+    SimpleUnaryImpl<I, O> {
   return (values, dtype, attrs) => {
     const newValues =
-        util.getTypedArrayFromDType(dtype as NumericDataType, values.length);
+        util.getArrayFromDType(dtype, values.length);
     for (let i = 0; i < values.length; ++i) {
       newValues[i] = op(values[i], attrs);
     }

tfjs-backend-cpu/src/utils/unary_types.ts

@@ -15,8 +15,12 @@
  * =============================================================================
  */
 
-import {DataType, NamedAttrMap, TypedArray} from '@tensorflow/tfjs-core';
+import {DataTypeFor, DataTypeMap, NamedAttrMap} from '@tensorflow/tfjs-core';
 
-export type SimpleUnaryOperation = (x: number, attrs?: NamedAttrMap) => number;
-export type SimpleUnaryImpl =
-    (values: TypedArray, dtype: DataType, attrs?: NamedAttrMap) => TypedArray;
+export type SimpleUnaryOperation<I extends number | string = number,
+    O extends number | string = number> = (x: I, attrs?: NamedAttrMap) => O;
+
+export type SimpleUnaryImpl<I extends number | string = number | string,
+  O extends number | string = number | string> =
+  (values: ArrayLike<I>, dtype: DataTypeFor<O>,
+   attrs?: NamedAttrMap) => DataTypeMap[DataTypeFor<O>];

tfjs-backend-cpu/src/utils/unary_utils.ts

@@ -15,10 +15,11 @@
  * =============================================================================
  */
 
-import {DataType, KernelFunc, TypedArray, UnaryInputs, util} from '@tensorflow/tfjs-core';
+import {backend_util, DataTypeFor, KernelFunc, UnaryInputs} from '@tensorflow/tfjs-core';
 
 import {MathBackendCPU} from '../backend_cpu';
 import {assertNotComplex} from '../cpu_util';
+import {createSimpleUnaryImpl} from './unary_impl';
 
 import {SimpleUnaryImpl, SimpleUnaryOperation} from './unary_types';
 
@@ -30,25 +31,14 @@ import {SimpleUnaryImpl, SimpleUnaryOperation} from './unary_types';
  *     result has the same dtype as the input. This is mainly used in certain
  *     kernels that return bool type, such as isFinite, isInf, etc.
  */
-export function unaryKernelFunc(
-    name: string, op: SimpleUnaryOperation, dtype?: DataType): KernelFunc {
-  return ({inputs, attrs, backend}) => {
-    const {x} = inputs as UnaryInputs;
-    assertNotComplex(x, name);
-    if (x.dtype === 'string' || dtype === 'string') {
-      throw new Error('unaryKernelFunc does not support string input/output');
-    }
+export function unaryKernelFunc<I extends number | string = number,
+  O extends number | string = number>(
+  name: string, op: SimpleUnaryOperation<I, O>,
+  dtype?: DataTypeFor<O>): KernelFunc {
 
-    const cpuBackend = backend as MathBackendCPU;
-    const values = cpuBackend.data.get(x.dataId).values as TypedArray;
-    const xSize = util.sizeFromShape(x.shape);
-    const $dtype = dtype || x.dtype;
-    const newValues = util.getArrayFromDType($dtype, xSize);
-    for (let i = 0; i < xSize; ++i) {
-      newValues[i] = op(values[i], attrs);
-    }
-    return cpuBackend.makeTensorInfo(x.shape, $dtype, newValues);
-  };
+  const impl = createSimpleUnaryImpl<I, O>(op);
+
+  return unaryKernelFuncFromImpl<I, O>(name, impl, dtype);
 }
 
 /**
@@ -60,19 +50,30 @@ export function unaryKernelFunc(
  *     result has the same dtype as the input. This is mainly used in certain
  *     kernels that return bool type, such as isFinite, isInf, etc.
  */
-export function unaryKernelFuncFromImpl(
-    name: string, unaryImpl: SimpleUnaryImpl, dtype?: DataType): KernelFunc {
+export function unaryKernelFuncFromImpl<I extends number | string = number,
+  O extends number | string = number>(
+  name: string, unaryImpl: SimpleUnaryImpl<I, O>,
+  dtype?: DataTypeFor<O>): KernelFunc {
+
   return ({inputs, attrs, backend}) => {
     const {x} = inputs as UnaryInputs;
     assertNotComplex(x, name);
-    if (x.dtype === 'string' || dtype === 'string') {
-      throw new Error('unaryKernelFunc does not support string input/output');
-    }
 
     const cpuBackend = backend as MathBackendCPU;
-    const values = cpuBackend.data.get(x.dataId).values as TypedArray;
-    const $dtype = dtype || x.dtype;
-    const newValues = unaryImpl(values, $dtype, attrs);
+    const values = cpuBackend.data.get(x.dataId).values;
+    let decoded: ArrayLike<I>;
+    if (x.dtype === 'string') {
+      if (!Array.isArray(values)) {
+        throw new Error('String tensor\'s value was not an instance of Array');
+      }
+      decoded = backend_util.fromUint8ToStringArray(values) as unknown as
+        ArrayLike<I>;
+    } else {
+      decoded = values as unknown as ArrayLike<I>;
+    }
+
+    const $dtype = dtype || x.dtype as DataTypeFor<O>;
+    const newValues = unaryImpl(decoded, $dtype, attrs);
     return cpuBackend.makeTensorInfo(x.shape, $dtype, newValues);
   };
 }

tfjs-backend-webgl/src/kernel_utils/kernel_funcs_utils.ts

@@ -15,7 +15,7 @@
  * =============================================================================
  */
 
-import {backend_util, BinaryInputs, DataType, env, KernelFunc, TypedArray, UnaryInputs, upcastType} from '@tensorflow/tfjs-core';
+import { backend_util, BinaryInputs, DataType, env, KernelFunc, TypedArray, UnaryInputs, upcastType} from '@tensorflow/tfjs-core';
 
 import {MathBackendWebGL} from '../backend_webgl';
 import {BinaryOpProgram} from '../binaryop_gpu';
@@ -36,7 +36,7 @@ type UnaryKernelFuncConfig = {
   opSnippet: string,
   packedOpSnippet?: string,
   cpuKernelImpl?: SimpleUnaryKernelImplCPU,
-  dtype?: DataType
+  dtype?: DataType,
 };
 
 /**

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

@@ -66,6 +66,7 @@ const {
   sparseReshapeImpl: sparseReshapeImplCPU,
   sparseSegmentReductionImpl: sparseSegmentReductionImplCPU,
   sqrtImpl: sqrtImplCPU,
+  staticRegexReplaceImpl: staticRegexReplaceImplCPU,
   stridedSliceImpl: stridedSliceImplCPU,
   stringNGramsImpl: stringNGramsImplCPU,
   stringSplitImpl: stringSplitImplCPU,
@@ -114,6 +115,7 @@ export {
   sparseReshapeImplCPU,
   sparseSegmentReductionImplCPU,
   sqrtImplCPU,
+  staticRegexReplaceImplCPU,
   stridedSliceImplCPU,
   stringNGramsImplCPU,
   stringSplitImplCPU,

tfjs-backend-webgl/src/kernels/StaticRegexReplace.ts

@@ -0,0 +1,47 @@
+/**
+ * @license
+ * Copyright 2023 Google LLC.
+ * 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 {backend_util, KernelConfig, KernelFunc, NamedAttrMap, StaticRegexReplace, StaticRegexReplaceAttrs, StaticRegexReplaceInputs, TensorInfo} from '@tensorflow/tfjs-core';
+import {MathBackendWebGL} from '../backend_webgl';
+import {staticRegexReplaceImplCPU} from '../kernel_utils/shared';
+
+export function staticRegexReplace(args: {
+  inputs: StaticRegexReplaceInputs,
+  backend: MathBackendWebGL,
+  attrs: StaticRegexReplaceAttrs,
+}): TensorInfo {
+  const {inputs, backend, attrs} = args;
+  const {x} = inputs;
+
+  if (x.dtype !== 'string') {
+    throw new Error('Input must be of datatype string');
+  }
+
+  const $x = backend.readSync(x.dataId) as Uint8Array[];
+
+  const stringInput = backend_util.fromUint8ToStringArray($x);
+  const output = staticRegexReplaceImplCPU(stringInput, 'string',
+                                           attrs as unknown as NamedAttrMap);
+
+  return backend.makeTensorInfo(x.shape, 'string', output);
+}
+
+export const staticRegexReplaceConfig: KernelConfig = {
+  kernelName: StaticRegexReplace,
+  backendName: 'webgl',
+  kernelFunc: staticRegexReplace as unknown as KernelFunc,
+};

tfjs-backend-webgl/src/register_all_kernels.ts

@@ -166,6 +166,7 @@ import {splitVConfig} from './kernels/SplitV';
 import {sqrtConfig} from './kernels/Sqrt';
 import {squareConfig} from './kernels/Square';
 import {squaredDifferenceConfig} from './kernels/SquaredDifference';
+import {staticRegexReplaceConfig} from './kernels/StaticRegexReplace';
 import {stepConfig} from './kernels/Step';
 import {stridedSliceConfig} from './kernels/StridedSlice';
 import {stringNGramsConfig} from './kernels/StringNGrams';
@@ -337,6 +338,7 @@ const kernelConfigs: KernelConfig[] = [
   sqrtConfig,
   squareConfig,
   squaredDifferenceConfig,
+  staticRegexReplaceConfig,
   stepConfig,
   stridedSliceConfig,
   stringNGramsConfig,