I am using Lux.jl to implement models whose weights are complex, and the output can be real or complex. If i keep the output real and the weights complex, I cannot compute gradients because the gradient container is determined by the type of the model output (this is what my understanding is). But for my use case I want real output, and I want the gradient container type to be determined by the superset of all weight types. For example if there are both real and complex weights, then the gradient container should be complex since that is the bigger set (I’m not sure how to express this correctly - my intuition is Real < Complex < Number). Where can I make changes to have this behaviour?