Amazing, I was looking to do something very similar for myself-- dig into pytorch library to see how everything was working.
Julia: https://juliadiff.org/ChainRulesCore.jl/stable/FAQ.html#Wher...? (bottom)
OCaml: https://ocaml.xyz/book/
The only language that you might actually use is Python.