fix wasm strided slice with new axis (#4235)

BUG BUG Co-authored-by: Ping Yu <4018+pyu10055@users.noreply.github.com>
tensorflow · Mar 2, 2021 · bf1a062 · bf1a062
1 parent 93e00bc
commit bf1a062
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/tfjs-backend-wasm/src/kernels/StridedSlice.ts b/tfjs-backend-wasm/src/kernels/StridedSlice.ts
@@ -111,7 +111,8 @@ export function stridedSlice(args: {
 
   const nonStrided = strides.every(v => v === 1);
   if (nonStrided) {
-    const xSliced = slice({inputs: {x}, attrs: {begin, size}, backend});
+    const xSliced = slice(
+        {inputs: {x: xReshaped}, attrs: {begin, size}, backend});
     backend.disposeData(xReshaped.dataId);
     const reshaped =
         reshape({inputs: {x: xSliced}, attrs: {shape: outShape}, backend});

diff --git a/tfjs-core/src/ops/strided_slice_test.ts b/tfjs-core/src/ops/strided_slice_test.ts
@@ -134,7 +134,7 @@ describeWithFlags('stridedSlice', ALL_ENVS, () => {
     expectArraysClose(await output.data(), [0, 2]);
   });
 
-  it('strided slice with several new axes', () => {
+  it('strided slice with several new axes', async () => {
     // Python slice code: t[1:2,tf.newaxis,0:3,tf.newaxis,2:5]
     const t = tf.zeros([2, 3, 4, 5]);
     const begin = [1, 0, 0, 0, 2];
@@ -147,6 +147,7 @@ describeWithFlags('stridedSlice', ALL_ENVS, () => {
     const output = tf.stridedSlice(
         t, begin, end, strides, beginMask, endMask, ellipsisMask, newAxisMask);
     expect(output.shape).toEqual([1, 1, 3, 1, 2, 5]);
+    expectArraysClose(await output.data(), new Array(30).fill(0));
   });
 
   it('strided slice with new axes and shrink axes', () => {