One Hot: Fixes CPU result dtype (#1110)

BUG: tensorflow/tfjs#435
tensorflow · Jun 17, 2018 · 833475d · 833475d
1 parent b5b5430
commit 833475d
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 1 deletion.
diff --git a/src/kernels/backend_cpu.ts b/src/kernels/backend_cpu.ts
@@ -1888,7 +1888,7 @@ export class MathBackendCPU implements KernelBackend {
     for (let event = 0; event < indices.size; ++event) {
       res[event * depth + indices.get(event)] = onValue;
     }
-    return ops.tensor2d(res, [indices.size, depth]);
+    return ops.tensor2d(res, [indices.size, depth], 'int32');
   }
 
   private broadcastedBinaryOp(

diff --git a/src/ops/array_ops_test.ts b/src/ops/array_ops_test.ts
@@ -2162,6 +2162,14 @@ describeWithFlags('oneHot', ALL_ENVS, () => {
     const indices = tf.tensor1d([0, 1], 'float32');
     expect(() => tf.oneHot(indices, 2)).toThrowError();
   });
+
+  it('check output dtype', () => {
+    const expectedType = 'int32';
+    const indices = tf.tensor1d([0, 1], 'int32');
+    const res = tf.oneHot(indices, 2);
+
+    expect(res.dtype).toEqual(expectedType);
+  });
 });
 
 describeWithFlags('linspace', ALL_ENVS, () => {