Open avik-pal opened 2 years ago
I think the poor broadcasting performance likely has to do with some missed chance to perform our memory optimization in the broadcast logic somewhere (i.e., we think it is unsafe to optimize in a place where it's actually safe to do so). I will take a look into this when I get a chance, but thanks for putting this together and providing a nice, realistic benchmark for performance going forward!
Also, minor nit, but ImmutableArray
will copy by itself if we can't optimize, so don't think copy
is necessary here:
make_immutable(x::AbstractArray) = ImmutableArray(copy(x))
I added it for ReshapedArray
doesn't seem to have a dispatch for that
Testing out the Immutable Arrays from https://github.com/JuliaLang/julia/pull/44381 with #7
TLDR: Performance is a slight pain (seems broadcasting) right now, but it is very straightforward to support these once the functionality is available in Base
EDIT: Code updated to work for Lux 0.4.*
Trial 1: From the Usage Example
Standard Abstract Arrays
Immutable Arrays
Trial 2: Only a Dense Layer
Standard Abstract Arrays
Immutable Arrays
Seems like there is a lot of time being spent on broadcasting the bias (seems like a problem with broadcasting in general)
Trial 3: No broadcasting
Standard Abstract Arrays
Immutable Arrays
Trial 4
Standard Abstract Arrays
Immutable Arrays
cc @ChrisRackauckas @ianatol @aviatesk