What is the most efficient way to automatically support custom data type?

I’m currently trying to support Paillier Encryption as the custom type of PyTorch. I already have supported basic addition, multiplication, and nn.Linear module. The examples can be found at 10.1. Paillier Encryption — AIJack documentation. The current implementation (https://github.com/Koukyosyumei/AIJack/blob/main/src/aijack/defense/paillier/torch_wrapper.py
) wraps each operation one by one manually, but I want to automatically support as many operations as possible like conv2d without manual implementation. Any idea would be helpful. I really appreciate any help you can provide.