reshape fix (#7678)

This commit is contained in:
Xueyun Zhu 2021-05-13 16:34:03 -07:00 committed by GitHub
parent 4b37901f10
commit 32d8278c2f
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 11 additions and 9 deletions

View file

@ -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] {

View file

@ -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],
},
];
}