Check out Jax: https://jax.readthedocs.io/en/latest/notebooks/autodiff_cook...

I remember seeing this one and an ocaml versions while back.

https://gist.github.com/ttesmer/948df432cf46ec6db8c1e83ab59b...

I wrote a purely functional AD library in Mercury [0], which adapts a general approach from [1]. I believe that Owl provides a similar approach [2].

[0] https://github.com/mclements/mercury-ad

[1] https://github.com/qobi/AD-Rosetta-Stone/

[2] https://github.com/owlbarn/owl/tree/main/src/base/algodiff