I’d probably try the JIT (but I don’t have a whole lot of experience with it). As far as I understand, the goal is to support most reasonable things.
These might have been previously supported, originally, a lot more autograd stuff was going on in Python (and the genuine Python autograd.Function derivatives still have it).
If the JIT doesn’t cut the mustard for your use-case you look into whether support is desirable. The Contributing document 4 rightly emphasizes the importance of discussing the feature with the core-devs first and then start the implementation.
Best regards
Thomas