Speeding up my logsumexp function

It would probably be easy to make vreduce work, by defining a gradient as in this PR using perhaps Tracker’s as the model.

To make @tullio work you’d have to solve that issue, the easy part is Base.gensym(ex::Expr) = gensym(string(ex)) but I didn’t see how to make @grad deal with callable structs. (The macro wraps the forward & backward functions in a struct Eval so that there’s somewhere to attach gradient definitions.)

Otherwise you may need to go back to write things as broadcasting instead. Or tell Turing to use Zygote not ReverseDiff.

1 Like