Fix SparseFillEmptyRows output shape bug (tensorflow#5020)

Lay2000 · May 3, 2021 · 1dbe52b · 1dbe52b
1 parent 9370884
commit 1dbe52b
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 2 deletions.
diff --git a/tfjs-backend-cpu/src/kernels/SparseFillEmptyRows_impl.ts b/tfjs-backend-cpu/src/kernels/SparseFillEmptyRows_impl.ts
@@ -130,7 +130,7 @@ export function sparseFillEmptyRowsImpl(
       }
     }
     return [
-      outputIndices, [indicesCount, rank], outputValues, emptyRowIndicator,
+      outputIndices, [fullIndicesCount, rank], outputValues, emptyRowIndicator,
       reverseIndexMap
     ];
   }

diff --git a/tfjs-core/src/ops/sparse/sparse_fill_empty_rows.ts b/tfjs-core/src/ops/sparse/sparse_fill_empty_rows.ts
@@ -62,7 +62,7 @@ import {op} from '../operation';
  * result['outputIndices'].print(); // [[0, 0], [1, 0], [1, 3], [1, 4],
  *                                  //  [2, 0], [3, 2], [3, 3], [4, 0]]
  * result['outputValues'].print(); // [0, 10, 13, 14,-1, 32, 33, -1]
- * result['emptyRowIndicator'].print(); // [0, 0, 1, 0, 1]
+ * result['emptyRowIndicator'].print(); // [false, false, true, false, true]
  * result['reverseIndexMap'].print(); // [0, 1, 2, 3, 5, 6]
  * ```
  * @param indices: 2-D. the indices of the sparse tensor.

diff --git a/tfjs-core/src/ops/sparse/sparse_fill_empty_rows_test.ts b/tfjs-core/src/ops/sparse/sparse_fill_empty_rows_test.ts
@@ -100,6 +100,35 @@ describeWithFlags('sparseFillEmptyRows', ALL_ENVS, () => {
     expectArraysClose(await result.reverseIndexMap.data(), []);
   });
 
+  it('check output metadata', async () => {
+    const sparseTensor = {
+      ind: tf.tensor2d(
+          [[0, 0], [1, 0], [1, 3], [1, 4], [3, 2], [3, 3]], [6, 2], 'int32'),
+      val: tf.tensor1d([0, 10, 13, 14, 32, 33], 'float32'),
+      shape: [5, 6],
+    };
+    const result = tf.sparse.sparseFillEmptyRows(
+        sparseTensor.ind, sparseTensor.val, sparseTensor.shape, -1);
+
+    expectArraysClose(
+        await result.outputIndices.data(),
+        [[0, 0], [1, 0], [1, 3], [1, 4], [2, 0], [3, 2], [3, 3], [4, 0]]);
+    expectArraysClose(
+        await result.outputValues.data(), [0, 10, 13, 14, -1, 32, 33, -1]);
+    expectArraysClose(await result.emptyRowIndicator.data(), [0, 0, 1, 0, 1]);
+    expectArraysClose(await result.reverseIndexMap.data(), [0, 1, 2, 3, 5, 6]);
+
+    expect(result.outputIndices.shape).toEqual([8, 2]);
+    expect(result.outputValues.shape).toEqual([8]);
+    expect(result.emptyRowIndicator.shape).toEqual([5]);
+    expect(result.reverseIndexMap.shape).toEqual([6]);
+
+    expect(result.outputIndices.dtype).toEqual(sparseTensor.ind.dtype);
+    expect(result.outputValues.dtype).toEqual(sparseTensor.val.dtype);
+    expect(result.emptyRowIndicator.dtype).toEqual('bool');
+    expect(result.reverseIndexMap.dtype).toEqual(sparseTensor.ind.dtype);
+  });
+
   it('throw error if dense rows is empty and indices is not', async () => {
     const sparseTensor = {
       ind: tf.tensor2d([[0, 0]], [1, 2], 'int32'),