I have reverse-mode (purely functional reverse mode at that!) sitting in a branch, and will get this going at some point soon. Even more fun will be compilation down to XLA, like JAX does in Python.