Matrix-based predictions and complex variables

As long as all the complex-valued operations happen inside Op that takes in a real and outputs a real, it shouldn’t be a problem. Be aware that you will have to define the gradients manually and implement those as well (the gradients should also be \mathbb R \mapsto \mathbb R)