Incompatibility of compute_dtype with complex-valued inputs

[jhoydis]

Hi,

In #19872, you introduced the possibility for layers with complex-valued inputs.

It then seems that this statement of the API Documentation is now wrong:

When I feed a complex-valued input tensor into a layer (as in this unit test), it is not cast to the compute_dtype, but rather kept as it is. I would somehow expect that the compute_dtype becomes complex in this case as well.

[mehtamansi29]

Hi @jhoydis -

Thanks for reporting this issue. Currently compute_dtype on layers automatically cast inputs to the dtype. But for complex input is not casting the same dtype. As complex values need to handle differently for casting compute_dtype.

We will dig into this issue more and update soon. Thanks..!!