[js/webgpu] minor fixes to make tinyllama work (#19564)

This commit is contained in:
Guenther Schmuelling 2024-02-23 15:45:30 -08:00 committed by GitHub
parent aec2389ad0
commit bb43a0f133
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 4 additions and 2 deletions

View file

@ -154,7 +154,9 @@ const createConcatProgramInfo = (inputs: readonly TensorView[], axis: number): P
export const concat = (context: ComputeContext, attributes: ConcatAttributes): void => {
validateInputs(context.inputs);
context.compute(createConcatProgramInfo(context.inputs, attributes.axis));
// 0 length tensors are valid for concat, remove them
const nonEmptyInputs = context.inputs.filter(input => ShapeUtil.size(input.dims) > 0);
context.compute(createConcatProgramInfo(nonEmptyInputs, attributes.axis), {inputs: nonEmptyInputs});
};
export const parseConcatAttributes = (attributes: Record<string, unknown>): ConcatAttributes =>

View file

@ -55,7 +55,7 @@ const createGatherProgramInfo = (inputs: readonly TensorView[], attributes: Gath
if (idx${x} < 0) {
idx${x} = idx${x} + uniforms.axisDimLimit;
}
var dataIndices${x} = ${data.type.indices}(0);
var dataIndices${x} : ${data.type.indices};
`;
for (let i = 0, j = 0; i < inputRank; i++) {
if (i === axis) {