Skip to content

Decrease load time by moving everything except stubs out to ChainRules.jl? #413

@oxinabox

Description

@oxinabox

This issue is being made after the fact, but so we have a record.
It has been proposed a few times moving out everything except the rrule, frule stubs,
the ProjectTo type, and the tangent types out to ChainRules.jl
to decrease load time.
So getting rid of all the linear algebra overloads for tangent types, and all the actual defintiions of ProjectTos.
e.g. mentioned here #340

This would not break any real code trying to do this with AD,
since any AD would also load ChainRules.jl where these would be defined.

After running a bunch of timing experiments, we have concluded not to do this.
The improvement boils down to a 0.025-0.055.
Which is a substantial portion of the load time.
But is still a pretty minuscule amount of time.

The downside of removing it is that it becomes much more annoying to try things out in the REPL for debugging your rrules.
Not only do you need to load ChainRulesCore, which is already a dependency of your package.
You also will need to load ChainRules.jl which isn't.
It seems like the time spent fussing around with environments will dwarf the time savings.
Keeping this stuff there means you have everything you need to define and use a rule, in one place.
We can revisit this for v2.0.

Data:

It can be noted that the plot breaks into two groups.
The difference is as to if this is the first package loaded after Revise.jl (which many people have in their startup.jl)
It can be seen that vs the overhead of initially loading after Revise, this is nothing.
image

Spreedsheet

Timing scripts:

using Pkg
Pkg.activate(temp=true)
pkg"add ChainRulesCore#master"
versioninfo()
pkg"st"
@time using ChainRulesCore
using Pkg
Pkg.activate(temp=true)
pkg"add ChainRulesCore#ox/mini"
versioninfo()
pkg"st"
@time using ChainRulesCore

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions