GEMM operator calculates the c output shape incorrectly when input a is transpose?

[kandrio]

Hey everyone!

I believe I have found a bug in the GEMM operator.

To the best of my knowledge, the output shape of the c StorageView in the GEMM operator should always be: {m, n}.

However, Gemm::compute() calcultes the shape of c incorrectly if a is transpose (a has shape {k, m} instead of {m, k}).

Take a look at this piece of code:

CTranslate2/src/ops/gemm.cc

Lines 84 to 88 in 4f8a4f3

	{
	Shape output_shape(a.shape());
	output_shape[output_shape.size() - 1] = n;
	c.resize(std::move(output_shape));
	}

• Shape output_shape(a.shape());: sets output_shape to: {k, m}
• output_shape[output_shape.size() - 1] = n;: sets output_shape to: {k, n}
• c.resize(std::move(output_shape));: resizes c to {k, n}.

Am I missing something here?

Bare in mind that there are no unit tests to catch this currently.

[minhthuc2502]

Hello,
This is only the high level of gemm ops. The shape of a is always {m, k}. The private member of Gemm class is _trans_a which is used in the low level of gemm ops for the later computation in src/cpu/primitives.cc

[kandrio]

Hi @minhthuc2502, thanks for answering.

The shape of a is always {m, k}.

I would argue that this is not true. Take a look at this line. Notice how k gets its value from the 1st dimension of a if _trans_a==true. So, in that case, the shape of a is {k, m} (as expected).

The private member of Gemm class is _trans_a which is used in the low level of gemm ops for the later computation in src/cpu/primitives.cc

I don't doubt that, but I still think that c gets a completely wrong shape ({k, n} and k is completely out of place here) in the case that _trans_a==true.

What do you think?

[minhthuc2502]

You are right, I think we never handle the case where a is transpose. I will do the fix asap. Thank you for your report.

[minhthuc2502]

I have to reconfirm that your initial idea is correct. However, the problem is gemm operation in this project is used for a whose shape is 3d like batch_size x m x k. If the shape of a is batch_size x k x m, a transpose would be m x (batch_size * k), and it would not be multiple with the b (shape k x n). In the operation matmul here, it does well the case where a is transpose. I will close this issue, if you have any issue with that, feel free to open new issue.