diff --git a/js/web/lib/onnxjs/backends/webgl/ops/reshape-packed.ts b/js/web/lib/onnxjs/backends/webgl/ops/reshape-packed.ts index 11f287de61..8e023536e3 100644 --- a/js/web/lib/onnxjs/backends/webgl/ops/reshape-packed.ts +++ b/js/web/lib/onnxjs/backends/webgl/ops/reshape-packed.ts @@ -42,13 +42,11 @@ export class WebGLReshapePacked extends Reshape implements WebGLOperator { handler.createTextureLayoutFromShape(inputShape3D, 4, inputShape3D, {isPacked: true, reverseWH: true}); } - this.outputShape = ShapeUtil.calculateReshapedDims(originInputShape, inputs[1].integerData); - const squeezedOutputShape = processDims3D(this.outputShape); + const outputShape = ShapeUtil.calculateReshapedDims(originInputShape, inputs[1].integerData); + const squeezedOutputShape = processDims3D(outputShape); - const outputLayout = handler.createTextureLayoutFromShape( + this.outputLayout = handler.createTextureLayoutFromShape( squeezedOutputShape, 4, squeezedOutputShape, {isPacked: true, reverseWH: true}); - this.originalOutputLayout = - handler.createTextureLayoutFromShape(this.outputShape, 4, this.outputShape, {isPacked: true, reverseWH: true}); let mainLoop = ''; for (let i = 0; i < 4; i++) { @@ -106,7 +104,7 @@ export class WebGLReshapePacked extends Reshape implements WebGLOperator { return { inputLayouts: [inputLayout], - outputLayout, + outputLayout: this.outputLayout, samplers: ['A'], shaderSource, hasMain: true, @@ -117,7 +115,7 @@ export class WebGLReshapePacked extends Reshape implements WebGLOperator { createRunData(handler: WebGLInferenceHandler, programInfo: ProgramInfo, inputs: Tensor[]): RunData { const inputTDs = [handler.getOrCreateTextureData(inputs[0], handler.getOrCreateTextureLayout(inputs[0], 1, false, [], false))]; - let outputLayout = this.originalOutputLayout; + let outputLayout = this.outputLayout; if (outputLayout === undefined) { const originInputShape = inputs[0].dims; const outputShape = ShapeUtil.calculateReshapedDims(originInputShape, inputs[1].integerData); @@ -131,8 +129,7 @@ export class WebGLReshapePacked extends Reshape implements WebGLOperator { uniformData: {} }; } - protected outputShape: readonly number[]; - private originalOutputLayout: TextureLayout; + private outputLayout: TextureLayout; } function processDims3D(shape: readonly number[]|readonly number[]|Tensor.IntegerType): [number, number, number] { diff --git a/js/web/test/unittests/backends/webgl/test-reshape-packed.ts b/js/web/test/unittests/backends/webgl/test-reshape-packed.ts index e89ee047ff..cfd6e3d459 100644 --- a/js/web/test/unittests/backends/webgl/test-reshape-packed.ts +++ b/js/web/test/unittests/backends/webgl/test-reshape-packed.ts @@ -90,6 +90,11 @@ function getTestData(): TestData[] { inputShape: [2, 2, 2, 4], outputShape: [2, 1, 4, 4], }, + { + elementCount: 18432, + inputShape: [512, 36], + outputShape: [512, 36, 1, 1], + }, ]; }