fix conjugate-transpose for matrices of certain sizes (issue #19200)

[OmriSteiner]

When using tf.linalg.matrix_transpose or tf.linalg.adjoint to conjugate-transpose a complex tensor, the functions might fail to conjugate the tensor under certain circumstances.

Given a complex matrix with m rows and n columns, if the following conditions are met:

1. The operation runs on GPU.
2. (m > 96 and n < 16) or (n > 96 and n < 16).
   Then the matrix will not be conjugated (although it will be transposed).

Code snippet which reproduces the bug (before patch):

import tensorflow as tf

with tf.device("GPU"):
    A = tf.complex(0., tf.ones(shape=(97,2)))
    Ah = tf.linalg.matrix_transpose(A, conjugate=True)

A[0][0] == Ah[0][0] ### True - but should be False

[google-cla]

Thanks for your pull request. It looks like this may be your first contribution to a Google open source project (if not, look below for help). Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

📝 Please visit https://cla.developers.google.com/ to sign.

Once you've signed (or fixed any issues), please reply here with @googlebot I signed it! and we'll verify it.

---

What to do if you already signed the CLA

Individual signers

• It's possible we don't have your GitHub username or you're using a different email address on your commit. Check your existing CLA data and verify that your email is set on your git commits.

Corporate signers

• Your company has a Point of Contact who decides which employees are authorized to participate. Ask your POC to be added to the group of authorized contributors. If you don't know who your Point of Contact is, direct the Google project maintainer to go/cla#troubleshoot (Public version).
• The email used to register you as an authorized contributor must be the email used for the Git commit. Check your existing CLA data and verify that your email is set on your git commits.
• The email used to register you as an authorized contributor must also be attached to your GitHub account.

ℹ️ Googlers: Go here for more info.

[OmriSteiner]

@googlebot I signed it!

[google-cla]

Thanks for your pull request. It looks like this may be your first contribution to a Google open source project (if not, look below for help). Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

📝 Please visit https://cla.developers.google.com/ to sign.

Once you've signed (or fixed any issues), please reply here with @googlebot I signed it! and we'll verify it.

---

What to do if you already signed the CLA

Individual signers

• It's possible we don't have your GitHub username or you're using a different email address on your commit. Check your existing CLA data and verify that your email is set on your git commits.

Corporate signers

• Your company has a Point of Contact who decides which employees are authorized to participate. Ask your POC to be added to the group of authorized contributors. If you don't know who your Point of Contact is, direct the Google project maintainer to go/cla#troubleshoot (Public version).
• The email used to register you as an authorized contributor must be the email used for the Git commit. Check your existing CLA data and verify that your email is set on your git commits.
• The email used to register you as an authorized contributor must also be attached to your GitHub account.

ℹ️ Googlers: Go here for more info.

[gbaned]

@OmriSteiner Can you please sign CLA. Thanks!

[OmriSteiner]

@gbaned My employer signed the CLA and added me as an authorized contributor, but apparently it takes several days for approval?

[OmriSteiner]

@googlebot I signed it!

[allenlavoie]

Can you add a unit test? Probably this file: https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/kernel_tests/array_ops_test.py

[allenlavoie]

Is this necessary for correctness? Would we run into similar issues with complex<float16> and TransposeElemType<4>?

[OmriSteiner]

Yes, this is necessary. Instead of simply conjugating, we call maybe_conj, which is specialized for complex types in GPU (float2 and double2).

I thought about TransposeElemType<4> when fixing this, and came to the conclusion that tensorflow doesn't have a complex type which uses float16. std::complex<float16> doesn't exist, as far as I can tell.

To be honest, I'm not entirely sure why TransposeElemType is there in the first place. If I had to make a guess, I would say that it's in order to make this code work for non-basic types (i.e not tensorflow builtin numeric types). However, this is only used in a very specific place, so it would fail to achieve that.

Thanks for the comment, I found out I didn't fix the issue for complex128, but it's sorted now.

[allenlavoie]

Oh, right, we just have DT_COMPLEX64 and DT_COMPLEX128. Sounds OK then. If you want to get rid of TransposeElemType that's fine too (we might find a failing test).

I do think it's another good reason to have a unit test. If someone comes along and adds a smaller complex type it'd make it at least possible to find and fix this issue.

[OmriSteiner]

Yeah, I could write a unit-test for this.

[OmriSteiner]

@allenlavoie I've expanded the test to also test complex128.