[Core] Fix bug with pointwise conv2d with packing

tfjs-core/src/backends/webgl/backend_webgl.ts

@@ -1923,8 +1923,11 @@ export class MathBackendWebGL extends KernelBackend {
     const targetShape = isChannelsLast ?
         xShape[0] * xShape[1] * (xShape[2] + 1) :
         xShape[0] * xShape[2] * (xShape[3] + 1);
-    const xReshaped = this.reshape(x, [1, targetShape, convInfo.inChannels]);
-
+    const xReshaped: TensorInfo = {
+      dataId: x.dataId,
+      shape: [1, targetShape, convInfo.inChannels],
+      dtype: x.dtype
+    };
     // xTexData.shape gets referenced from GPGPUBinary.inShapeInfos.
     // Decrementing row count, after batchMatMul->...->compileProgram leads to
     // invalid row count within the reference in GPGPUBinary.inShapeInfos.
@@ -1961,7 +1964,9 @@ export class MathBackendWebGL extends KernelBackend {
     // Set the output shape - there is no need for expensive reshape as data
     // layout is already correct.
     pointwiseConvTexData.shape = convInfo.outShape;
-    return this.reshape(pointwiseConv, convInfo.outShape);
+    return ENGINE.makeTensorFromDataId(
+               pointwiseConv.dataId, convInfo.outShape, pointwiseConv.dtype) as
+        Tensor4D;
   }
 
   private conv2dWithIm2Row(
@@ -2651,7 +2656,9 @@ export class MathBackendWebGL extends KernelBackend {
 
     if (!env().getBool('WEBGL_LAZILY_UNPACK') && outData.isPacked &&
         preventEagerUnpackingOfOutput === false) {
-      return this.unpackTensor(output);
+      const unpacked = this.unpackTensor(output);
+      this.disposeData(output.dataId);
+      return unpacked;
     }
     return output;
   }

tfjs-core/src/backends/webgl/webgl_ops_test.ts

@@ -17,7 +17,7 @@
 
 import * as tf from '../../index';
 import {describeWithFlags} from '../../jasmine_util';
-import {Tensor2D} from '../../tensor';
+import {Tensor2D, Tensor3D, Tensor4D} from '../../tensor';
 import {expectArraysClose, expectArraysEqual} from '../../test_util';
 import {Rank} from '../../types';
 
@@ -741,6 +741,32 @@ describeWithFlags('slice a packed texture', WEBGL_ENVS, () => {
   });
 });
 
+// TODO(annyuan, smilkov) Make this test pass for WEBGL_PACK=false.
+describeWithFlags('pointwise conv2d packed', WEBGL_ENVS, () => {
+  beforeAll(() => {
+    tf.env().set('WEBGL_SIZE_UPLOAD_UNIFORM', 0);
+  });
+
+  it('pointwise conv2d optimization with odd input size', async () => {
+    // We do special optimization in the webl backend which avoids an expensive
+    // reshape, when the following 3 conditions are met:
+    // 1) the input width/height is odd-shaped.
+    // 2) the input is already packed.
+    // 3) the filter size is 1x1, i.e. pointwise.
+    const inChannels = 1;
+    const outChannels = 2;
+    const oddInputSize = 3;
+    const x: Tensor3D = tf.ones([oddInputSize, oddInputSize, inChannels]);
+    const xPacked = tf.relu(x);
+    const pointwiseFilter: Tensor4D = tf.ones([1, 1, inChannels, outChannels]);
+    const strides = 1;
+    const pad = 'same';
+    const c = tf.conv2d(xPacked, pointwiseFilter, strides, pad);
+    expectArraysClose(
+        await c.data(), [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]);
+  });
+});
+
 describeWithFlags('relu', WEBGL_ENVS, () => {
   it('works with squarification for prime number length vector', async () => {
     const maxTextureSize = tf.env().getNumber('WEBGL_MAX_TEXTURE_SIZE');