Add Unique op

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

@@ -0,0 +1,46 @@
+/**
+ * @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, TensorInfo, Unique, UniqueAttrs, UniqueInputs} from '@tensorflow/tfjs-core';
+
+import {MathBackendCPU} from '../backend_cpu';
+import {assertNotComplex} from '../cpu_util';
+
+import {uniqueImpl} from './Unique_impl';
+
+export function unique(
+    args: {inputs: UniqueInputs, attrs: UniqueAttrs, backend: MathBackendCPU}):
+    TensorInfo[] {
+  const {inputs, attrs, backend} = args;
+  const {axis} = attrs;
+  const {x} = inputs;
+  assertNotComplex(x, 'unique');
+
+  const values = backend.data.get(x.dataId).values;
+  const {outputValues, outputShape, indices} =
+      uniqueImpl(values, axis, x.shape, x.dtype);
+  return [
+    backend.makeTensorInfo(outputShape, x.dtype, outputValues),
+    backend.makeTensorInfo([indices.length], 'int32', indices),
+  ];
+}
+
+export const uniqueConfig: KernelConfig = {
+  kernelName: Unique,
+  backendName: 'cpu',
+  kernelFunc: unique as {} as KernelFunc,
+};

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

@@ -0,0 +1,156 @@
+/**
+ * @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 {BackendValues, DataType, TensorBuffer, TypedArray, util} from '@tensorflow/tfjs-core';
+
+export function uniqueImpl(
+    values: BackendValues, axis: number, shape: number[], dtype: DataType): {
+  outputValues: BackendValues,
+  outputShape: number[],
+  indices: BackendValues
+} {
+  // Normalize and validate axis.
+  const $axis = util.parseAxisParam(axis, shape)[0];
+
+  // Calculate the new shape that is suitable for extracting data along the
+  // given axis.
+  //
+  // The rank is 3.
+  // The size of the 1st dimension is the size of all the axes < the given axis.
+  // The size of the 2nd dimension is the same as the size of the given axis.
+  // The size of the 3rd dimension is the size of all the axes > the given axis.
+  //
+  // For example, for a 4D tensor with shape=[2, 3, 5, 4] and axis=2, the
+  // newShape would be: [2*3, 5, 4].
+  //
+  // Note that this is not the final output shape. This will be the shape for an
+  // intermediate TensorBuffer (see inputBuffer below) to allow us to extract
+  // values along the given axis. To demonstrate how it works, consider the
+  // following example:
+  //
+  // Input: a 3D tensor, with shape [1, 2, 3]
+  // [
+  //   [
+  //      [1,2,3],
+  //      [4,5,6]
+  //   ]
+  // ]
+  // Axis: 2 (the last axis).
+  // Along axis 2, we expect to extract 3 tensors: [1,4], [2,5], [3,6].
+  //
+  // For this example, newShape would be: [2, 3, 1], where 2 is calculated from
+  // 1*2. The re-shaped data would look like:
+  //
+  // [
+  //   [
+  //     [1], [2], [3]
+  //   ],
+  //   [
+  //     [4], [5], [6]
+  //   ]
+  // ]
+  //
+  // Then, we can construct a 3-level nested loop by the following dimension
+  // order to extract the values along the axis (dimension1):
+  // i: dimension1       // 0,1,2 (newShape[1])
+  //   m: dimension0     // 0,1   (newShape[0])
+  //     n: dimension2   // 0     (newShape[2])
+  //
+  //                       m, i, n
+  //                      ---------
+  // Iteration 0: data at [0, 0, 0] => "1"
+  // Iteration 1: data at [1, 0, 0] => "4"
+  // We got [1,4].
+  // Iteration 2: data at [0, 1, 0] => "2"
+  // Iteration 3: data at [1, 1, 0] => "5"
+  // We got [2,5].
+  // Iteration 4: data at [0, 2, 0] => "3"
+  // Iteration 5: data at [1, 2, 0] => "6"
+  // We got [3,6].
+  const newShape = [1, shape[0], 1];
+  for (let i = 0; i < $axis; i++) {
+    newShape[0] *= shape[i];
+  }
+  newShape[1] = shape[$axis];
+  for (let i = $axis + 1; i < shape.length; i++) {
+    newShape[2] *= shape[i];
+  }
+
+  // A map from unique elements (their string representations) to their values
+  // in "indices" (below).
+  const uniqueElements: {[key: string]: number} = {};
+  // The indices of each unique element in the original tensor along the given
+  // axis. It is 1D and has the same size as the given axis.
+  const indices = new Int32Array(shape[$axis]);
+  // Create a buffer so we can easily extract value at a given location.
+  const inputBuffer = new TensorBuffer(newShape, dtype, values as TypedArray);
+  // The indices along the given axis that have unique elements. This is a
+  // de-duped version of "indices" above.
+  const uniqueIndices: number[] = [];
+  const is1DTensor = newShape[0] === 1 && newShape[2] === 1;
+  for (let i = 0; i < shape[$axis]; i++) {
+    // Extract values along the axis.
+    let element: string;
+    if (is1DTensor) {
+      // Fast path for 1D tensor input.
+      element = values[i].toString();
+    } else {
+      const axisValues = [];
+      for (let m = 0; m < newShape[0]; m++) {
+        for (let n = 0; n < newShape[2]; n++) {
+          axisValues.push(inputBuffer.get(m, i, n));
+        }
+      }
+      element = axisValues.join(',');
+    }
+
+    // Dedup and update various indices.
+    if (uniqueElements[element] !== undefined) {
+      indices[i] = uniqueElements[element];
+    } else {
+      const uniqueIndex = Object.keys(uniqueElements).length;
+      uniqueElements[element] = uniqueIndex;
+      indices[i] = uniqueIndex;
+      uniqueIndices.push(i);
+    }
+  }
+
+  // Now we know where each of the unique elements are located along the axis
+  // (uniqueIndices). Extract them from input buffer and store them in the
+  // output buffer.
+  const outputTmpShape = newShape.slice();
+  outputTmpShape[1] = Object.keys(uniqueElements).length;
+  const outputBuffer = new TensorBuffer(outputTmpShape, dtype);
+  uniqueIndices.forEach((uniqueElementIndex, i) => {
+    for (let m = 0; m < newShape[0]; m++) {
+      for (let n = 0; n < newShape[2]; n++) {
+        outputBuffer.set(inputBuffer.get(m, uniqueElementIndex, n), m, i, n);
+      }
+    }
+  });
+
+  // The output shape can be calculated from the input shape with the size of
+  // the given axis replaced by the number of unique elements along that axis.
+  const outputShape = shape.slice();
+  outputShape[$axis] = outputTmpShape[1];
+
+  return {
+    outputValues: outputBuffer.values as BackendValues,
+    outputShape,
+    indices,
+  };
+}

tfjs-backend-cpu/src/register_all_kernels.ts

@@ -88,6 +88,7 @@ import {subConfig} from './kernels/Sub';
 import {tanConfig} from './kernels/Tan';
 import {tanhConfig} from './kernels/Tanh';
 import {transposeConfig} from './kernels/Transpose';
+import {uniqueConfig} from './kernels/Unique';
 
 // List all kernel configs here
 const kernelConfigs: KernelConfig[] = [
@@ -159,7 +160,8 @@ const kernelConfigs: KernelConfig[] = [
   subConfig,
   tanConfig,
   tanhConfig,
-  transposeConfig
+  transposeConfig,
+  uniqueConfig,
 ];
 
 for (const kernelConfig of kernelConfigs) {

tfjs-backend-cpu/src/shared.ts

@@ -18,3 +18,4 @@
 // Shared kernel impls for use in other backends.
 export {maxImpl} from './kernels/Max_impl';
 export {transposeImpl} from './kernels/Transpose_impl';
+export {uniqueImpl} from './kernels/Unique_impl';

tfjs-backend-webgl/src/backend_webgl.ts

@@ -2360,8 +2360,9 @@ export class MathBackendWebGL extends KernelBackend {
     return backend_util.linspaceImpl(start, stop, num);
   }
 
-  makeTensorInfo(shape: number[], dtype: DataType): TensorInfo {
-    const dataId = this.write(null /* values */, shape, dtype);
+  makeTensorInfo(shape: number[], dtype: DataType, values?: BackendValues):
+      TensorInfo {
+    const dataId = this.write(values, shape, dtype);
     this.texData.get(dataId).usage = null;
     return {dataId, shape, dtype};
   }

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

@@ -20,6 +20,10 @@
 // tslint:disable-next-line: no-imports-from-dist
 import * as shared from '@tensorflow/tfjs-backend-cpu/dist/shared';
 
-const {maxImpl: maxImplCPU, transposeImpl: transposeImplCPU} = shared;
+const {
+  maxImpl: maxImplCPU,
+  transposeImpl: transposeImplCPU,
+  uniqueImpl: uniqueImplCPU,
+} = shared;
 
-export {maxImplCPU, transposeImplCPU};
+export {maxImplCPU, transposeImplCPU, uniqueImplCPU};

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

@@ -0,0 +1,50 @@
+/**
+ * @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, TensorInfo, Unique, UniqueAttrs, UniqueInputs} from '@tensorflow/tfjs-core';
+
+import {MathBackendWebGL} from '../backend_webgl';
+import {uniqueImplCPU} from '../kernel_utils/shared';
+import {assertNotComplex} from '../webgl_util';
+
+export function unique(
+    args:
+        {inputs: UniqueInputs, attrs: UniqueAttrs, backend: MathBackendWebGL}):
+    TensorInfo[] {
+  const {inputs, attrs, backend} = args;
+  const {axis} = attrs;
+  const {x} = inputs;
+  assertNotComplex(x, 'unique');
+
+  // For now, always forward calculation to the CPU backend.
+  console.warn(
+      'WARNING: ',
+      'UI might be locked temporarily as data is being downloaded');
+  const values = backend.readSync(x.dataId);
+  const {outputValues, outputShape, indices} =
+      uniqueImplCPU(values, axis, x.shape, x.dtype);
+  return [
+    backend.makeTensorInfo(outputShape, x.dtype, outputValues),
+    backend.makeTensorInfo([indices.length], 'int32', indices),
+  ];
+}
+
+export const uniqueConfig: KernelConfig = {
+  kernelName: Unique,
+  backendName: 'webgl',
+  kernelFunc: unique as {} as KernelFunc,
+};

tfjs-backend-webgl/src/register_all_kernels.ts

@@ -39,6 +39,7 @@ import {squareConfig} from './kernels/Square';
 import {squaredDifferenceConfig} from './kernels/SquaredDifference';
 import {tanConfig} from './kernels/Tan';
 import {transposeConfig} from './kernels/Transpose';
+import {uniqueConfig} from './kernels/Unique';
 
 // List all kernel configs here
 const kernelConfigs: KernelConfig[] = [
@@ -64,7 +65,8 @@ const kernelConfigs: KernelConfig[] = [
   squareConfig,
   squaredDifferenceConfig,
   tanConfig,
-  transposeConfig
+  transposeConfig,
+  uniqueConfig,
 ];
 
 for (const kernelConfig of kernelConfigs) {

tfjs-converter/docs/supported_ops.md

@@ -166,6 +166,8 @@
 |Tensorflow Op Name|Tensorflow.js Op Name|
 |---|---|
 |TopKV2|topK|
+|Unique|unique|
+|UniqueV2|unique|
 |Not mapped|confusionMatrix|
 |Not mapped|topk|
 

tfjs-converter/metadata/kernel2op.json

@@ -505,6 +505,12 @@
   "TruncatedNormal": [
     "truncatedNormal"
   ],
+  "Unique": [
+    "unique"
+  ],
+  "UniqueV2": [
+    "unique"
+  ],
   "Unpack": [
     "unstack"
   ],

tfjs-converter/python/tensorflowjs/op_list/evaluation.json

@@ -21,5 +21,32 @@
         "type": "bool"
       }
     ]
+  },
+  {
+    "tfOpName": "Unique",
+    "category": "evaluation",
+    "inputs": [
+      {
+        "start": 0,
+        "name": "x",
+        "type": "tensor"
+      }
+    ]
+  },
+  {
+    "tfOpName": "UniqueV2",
+    "category": "evaluation",
+    "inputs": [
+      {
+        "start": 0,
+        "name": "x",
+        "type": "tensor"
+      },
+      {
+        "start": 1,
+        "name": "axis",
+        "type": "number"
+      }
+    ]
   }
-]
+]

tfjs-converter/src/operations/executors/evaluation_executor.ts

@@ -37,6 +37,18 @@ export const executeOp: InternalOpExecutor =
               const result = tfOps.topk(x, k, sorted);
               return [result.values, result.indices];
             }
+            case 'Unique': {
+              const x = getParamValue('x', node, tensorMap, context) as Tensor;
+              const result = tfOps.unique(x);
+              return [result.values, result.indices];
+            }
+            case 'UniqueV2': {
+              const x = getParamValue('x', node, tensorMap, context) as Tensor;
+              const axis =
+                  getParamValue('axis', node, tensorMap, context) as number;
+              const result = tfOps.unique(x, axis);
+              return [result.values, result.indices];
+            }
             default:
               throw TypeError(`Node type ${node.op} is not implemented`);
           }