search

JuliaDiff / TaylorDiff.jl

Taylor-mode automatic differentiation for higher-order derivatives
https://juliadiff.org/TaylorDiff.jl/
MIT License
67 stars 7 forks source link

[TTFX] Use Packages Extensions to drop Zygote and others when not used. #62

Open lrnv opened 10 months ago

lrnv commented 10 months ago

TaylorDiff is depending on a lot of stuff it should not, which tanks my TTFX.

As already noted in https://github.com/JuliaDiff/TaylorDiff.jl/issues/52#issuecomment-1707801903_ , It would be very beneficial to leverage package preferences to avoid, say, having Zygote and ReverseDiff as dependencies. Formally, the composability with the AD systems is not a core functionality of the package and could be loaded ONLY when those AD systems are loaded.

lrnv commented 10 months ago

Taking a second look at this, It appear to me that it is only the last two files derivatives.jl and chainrules.jl that are really problematic.

1) In derivatives.jl, there is a dependency on the (old) SlicedMap.jl, which brings a lot of stuff with it. Since it is only used for three calls to their map functions, maybe you could reconsider and ditch it ? If you are worried about your perfs on Zygote stuff, then read point 2)

2) In chainrules.jl, the code is there to support a usecase with Zygote and/or with other AD systems. This code should definitely go into one (or two?) package extensions files so that loading Zygote and Chainrules is not necessary if I do not use them. The special overloades from SlicedMap.jl that are there for Zygote perfs could also be part of this package extension.

Tell me what you think

zhujch1 commented 10 months ago

  1. The dependency SlicedMap.jl has been removed. We have introduced a more efficient implementation.

  2. We are loading Zygote.jl for solving some issues in it, we will consider integrating ZygoteRules.jl as a weak dependency in the future. We can't remove Chainrules.jl, because we need to use 'frule's.

lrnv commented 10 months ago

For SliceMap: Thanks for removing it !

For Zygote : by weak-dep you mean package extension, right ?