chore: add NNlib as a direct dep

LuxDL · Jul 22, 2024 · 99a9bcb · 99a9bcb · github-actions · Jul 22, 2024
1 parent d81e2e8
commit 99a9bcb
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 0 deletions.
diff --git a/Project.toml b/Project.toml
@@ -24,6 +24,7 @@ LuxDeviceUtils = "34f89e08-e1d5-43b4-8944-0b49ac560553"
 LuxLib = "82251201-b29d-42c6-8e01-566dec8acb11"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 Markdown = "d6f4376e-aef5-505a-96c1-9c027394607a"
+NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
 Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
 Preferences = "21216c6a-2e73-6563-6e65-726566657250"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
@@ -98,6 +99,7 @@ MPI = "0.20.19"
 MacroTools = "0.5.13"
 Markdown = "1.10"
 NCCL = "0.1.1"
+NNlib = "0.9.21"
 OneHotArrays = "0.2.5"
 Optimisers = "0.3"
 Pkg = "1.10"

diff --git a/src/Lux.jl b/src/Lux.jl
@@ -18,6 +18,7 @@ using GPUArraysCore: GPUArraysCore, @allowscalar
 using LossFunctions: LossFunctions
 using MacroTools: MacroTools, block, combinedef, splitdef
 using Markdown: @doc_str
+using NNlib: NNlib
 using Optimisers: Optimisers
 using Preferences: load_preference, has_preference
 using Random: Random, AbstractRNG
Benchmark suite	Current: `99a9bcb`	Previous: `f2df403`	Ratio
`Dense(2 => 2)/cpu/reverse/ReverseDiff (compiled)/(2, 128)`	`3626.875` ns	`3654.375` ns	`0.99`
`Dense(2 => 2)/cpu/reverse/Zygote/(2, 128)`	`6876.333333333333` ns	`7153.357142857143` ns	`0.96`
`Dense(2 => 2)/cpu/reverse/Tracker/(2, 128)`	`20839` ns	`20799` ns	`1.00`
`Dense(2 => 2)/cpu/reverse/ReverseDiff/(2, 128)`	`9610.2` ns	`9912.5` ns	`0.97`
`Dense(2 => 2)/cpu/reverse/Flux/(2, 128)`	`8876.75` ns	`8970.6` ns	`0.99`
`Dense(2 => 2)/cpu/reverse/SimpleChains/(2, 128)`	`4440.875` ns	`4464.625` ns	`0.99`
`Dense(2 => 2)/cpu/reverse/Enzyme/(2, 128)`	`1155.3286713286714` ns	`1166.0921985815603` ns	`0.99`
`Dense(2 => 2)/cpu/forward/NamedTuple/(2, 128)`	`1109.3312883435583` ns	`1102.0060975609756` ns	`1.01`
`Dense(2 => 2)/cpu/forward/ComponentArray/(2, 128)`	`1176.6592592592592` ns	`1190.7692307692307` ns	`0.99`
`Dense(2 => 2)/cpu/forward/Flux/(2, 128)`	`1779.4285714285713` ns	`1804.78` ns	`0.99`
`Dense(2 => 2)/cpu/forward/SimpleChains/(2, 128)`	`179.88716502115656` ns	`184.81372549019608` ns	`0.97`
`Dense(20 => 20)/cpu/reverse/ReverseDiff (compiled)/(20, 128)`	`17202` ns	`17232` ns	`1.00`
`Dense(20 => 20)/cpu/reverse/Zygote/(20, 128)`	`16711.5` ns	`16872` ns	`0.99`
`Dense(20 => 20)/cpu/reverse/Tracker/(20, 128)`	`36839` ns	`37210` ns	`0.99`
`Dense(20 => 20)/cpu/reverse/ReverseDiff/(20, 128)`	`28468.5` ns	`28413` ns	`1.00`
`Dense(20 => 20)/cpu/reverse/Flux/(20, 128)`	`19838` ns	`20037` ns	`0.99`
`Dense(20 => 20)/cpu/reverse/SimpleChains/(20, 128)`	`17232` ns	`17343` ns	`0.99`
`Dense(20 => 20)/cpu/reverse/Enzyme/(20, 128)`	`4326.714285714285` ns	`4329.571428571428` ns	`1.00`
`Dense(20 => 20)/cpu/forward/NamedTuple/(20, 128)`	`3862.375` ns	`3883.5` ns	`0.99`
`Dense(20 => 20)/cpu/forward/ComponentArray/(20, 128)`	`3962.375` ns	`3951.8125` ns	`1.00`
`Dense(20 => 20)/cpu/forward/Flux/(20, 128)`	`4906.428571428572` ns	`4982.142857142857` ns	`0.98`
`Dense(20 => 20)/cpu/forward/SimpleChains/(20, 128)`	`1657.1` ns	`1659.1` ns	`1.00`
`Conv((3, 3), 3 => 3)/cpu/reverse/ReverseDiff (compiled)/(64, 64, 3, 128)`	`38622798` ns	`40947012` ns	`0.94`
`Conv((3, 3), 3 => 3)/cpu/reverse/Zygote/(64, 64, 3, 128)`	`57557404.5` ns	`58034514.5` ns	`0.99`
`Conv((3, 3), 3 => 3)/cpu/reverse/Tracker/(64, 64, 3, 128)`	`76110491` ns	`80730603` ns	`0.94`
`Conv((3, 3), 3 => 3)/cpu/reverse/ReverseDiff/(64, 64, 3, 128)`	`88543825` ns	`92327301` ns	`0.96`
`Conv((3, 3), 3 => 3)/cpu/reverse/Flux/(64, 64, 3, 128)`	`72540385.5` ns	`78411107` ns	`0.93`
`Conv((3, 3), 3 => 3)/cpu/reverse/SimpleChains/(64, 64, 3, 128)`	`12016501` ns	`11940034.5` ns	`1.01`
`Conv((3, 3), 3 => 3)/cpu/reverse/Enzyme/(64, 64, 3, 128)`	`6948863` ns	`7040440` ns	`0.99`
`Conv((3, 3), 3 => 3)/cpu/forward/NamedTuple/(64, 64, 3, 128)`	`7126539` ns	`7193096.5` ns	`0.99`
`Conv((3, 3), 3 => 3)/cpu/forward/ComponentArray/(64, 64, 3, 128)`	`7061273.5` ns	`7117252` ns	`0.99`
`Conv((3, 3), 3 => 3)/cpu/forward/Flux/(64, 64, 3, 128)`	`9951243` ns	`11651567` ns	`0.85`
`Conv((3, 3), 3 => 3)/cpu/forward/SimpleChains/(64, 64, 3, 128)`	`6393013` ns	`6422658` ns	`1.00`
`vgg16/cpu/reverse/Zygote/(32, 32, 3, 16)`	`693297402` ns	`701030424` ns	`0.99`
`vgg16/cpu/reverse/Zygote/(32, 32, 3, 64)`	`2549860453` ns	`2590201983` ns	`0.98`
`vgg16/cpu/reverse/Zygote/(32, 32, 3, 2)`	`130918802.5` ns	`140190010.5` ns	`0.93`
`vgg16/cpu/reverse/Tracker/(32, 32, 3, 16)`	`808148673` ns	`876551202` ns	`0.92`
`vgg16/cpu/reverse/Tracker/(32, 32, 3, 64)`	`2909311814` ns	`3356633010` ns	`0.87`
`vgg16/cpu/reverse/Tracker/(32, 32, 3, 2)`	`197970238` ns	`197466482` ns	`1.00`
`vgg16/cpu/reverse/Flux/(32, 32, 3, 16)`	`686649416.5` ns	`692417063` ns	`0.99`
`vgg16/cpu/reverse/Flux/(32, 32, 3, 64)`	`2648852469.5` ns	`2718741579` ns	`0.97`
`vgg16/cpu/reverse/Flux/(32, 32, 3, 2)`	`126515102` ns	`140643582` ns	`0.90`
`vgg16/cpu/forward/NamedTuple/(32, 32, 3, 16)`	`174887586.5` ns	`175950163.5` ns	`0.99`
`vgg16/cpu/forward/NamedTuple/(32, 32, 3, 64)`	`650058280` ns	`652447286` ns	`1.00`
`vgg16/cpu/forward/NamedTuple/(32, 32, 3, 2)`	`45387771` ns	`42722902.5` ns	`1.06`
`vgg16/cpu/forward/ComponentArray/(32, 32, 3, 16)`	`165356332.5` ns	`165562400` ns	`1.00`
`vgg16/cpu/forward/ComponentArray/(32, 32, 3, 64)`	`644256648` ns	`646079758.5` ns	`1.00`
`vgg16/cpu/forward/ComponentArray/(32, 32, 3, 2)`	`29957627.5` ns	`30503986` ns	`0.98`
`vgg16/cpu/forward/Flux/(32, 32, 3, 16)`	`210061510.5` ns	`216087564.5` ns	`0.97`
`vgg16/cpu/forward/Flux/(32, 32, 3, 64)`	`735900009.5` ns	`758513421.5` ns	`0.97`
`vgg16/cpu/forward/Flux/(32, 32, 3, 2)`	`40030481` ns	`38353748.5` ns	`1.04`
`Conv((3, 3), 64 => 64)/cpu/reverse/ReverseDiff (compiled)/(64, 64, 64, 128)`	`1256415603.5` ns	`1219847712.5` ns	`1.03`
`Conv((3, 3), 64 => 64)/cpu/reverse/Zygote/(64, 64, 64, 128)`	`1876933218.5` ns	`1887996340` ns	`0.99`
`Conv((3, 3), 64 => 64)/cpu/reverse/Tracker/(64, 64, 64, 128)`	`2407171626` ns	`2447955597` ns	`0.98`
`Conv((3, 3), 64 => 64)/cpu/reverse/ReverseDiff/(64, 64, 64, 128)`	`2559559973` ns	`2626651116` ns	`0.97`
`Conv((3, 3), 64 => 64)/cpu/reverse/Flux/(64, 64, 64, 128)`	`1858731782` ns	`1934007715` ns	`0.96`
`Conv((3, 3), 64 => 64)/cpu/reverse/Enzyme/(64, 64, 64, 128)`	`319059042` ns	`335642586.5` ns	`0.95`
`Conv((3, 3), 64 => 64)/cpu/forward/NamedTuple/(64, 64, 64, 128)`	`322002092` ns	`332570295` ns	`0.97`
`Conv((3, 3), 64 => 64)/cpu/forward/ComponentArray/(64, 64, 64, 128)`	`321704747` ns	`329527060` ns	`0.98`
`Conv((3, 3), 64 => 64)/cpu/forward/Flux/(64, 64, 64, 128)`	`404856363` ns	`374730028.5` ns	`1.08`
`Conv((3, 3), 1 => 1)/cpu/reverse/ReverseDiff (compiled)/(64, 64, 1, 128)`	`11746084` ns	`11990295.5` ns	`0.98`
`Conv((3, 3), 1 => 1)/cpu/reverse/Zygote/(64, 64, 1, 128)`	`17855711` ns	`18049600` ns	`0.99`
`Conv((3, 3), 1 => 1)/cpu/reverse/Tracker/(64, 64, 1, 128)`	`19121825` ns	`19367354` ns	`0.99`
`Conv((3, 3), 1 => 1)/cpu/reverse/ReverseDiff/(64, 64, 1, 128)`	`23754263` ns	`24028261` ns	`0.99`
`Conv((3, 3), 1 => 1)/cpu/reverse/Flux/(64, 64, 1, 128)`	`17804118` ns	`18064582` ns	`0.99`
`Conv((3, 3), 1 => 1)/cpu/reverse/SimpleChains/(64, 64, 1, 128)`	`1167930` ns	`1230918` ns	`0.95`
`Conv((3, 3), 1 => 1)/cpu/reverse/Enzyme/(64, 64, 1, 128)`	`2066844` ns	`2101198.5` ns	`0.98`
`Conv((3, 3), 1 => 1)/cpu/forward/NamedTuple/(64, 64, 1, 128)`	`2077422` ns	`2102362` ns	`0.99`
`Conv((3, 3), 1 => 1)/cpu/forward/ComponentArray/(64, 64, 1, 128)`	`2077595` ns	`2105777` ns	`0.99`
`Conv((3, 3), 1 => 1)/cpu/forward/Flux/(64, 64, 1, 128)`	`2068086` ns	`2098950` ns	`0.99`
`Conv((3, 3), 1 => 1)/cpu/forward/SimpleChains/(64, 64, 1, 128)`	`197462` ns	`214702` ns	`0.92`
`Dense(200 => 200)/cpu/reverse/ReverseDiff (compiled)/(200, 128)`	`291940` ns	`296645.5` ns	`0.98`
`Dense(200 => 200)/cpu/reverse/Zygote/(200, 128)`	`264327` ns	`266860` ns	`0.99`
`Dense(200 => 200)/cpu/reverse/Tracker/(200, 128)`	`365018` ns	`369252` ns	`0.99`
`Dense(200 => 200)/cpu/reverse/ReverseDiff/(200, 128)`	`406756` ns	`412393` ns	`0.99`
`Dense(200 => 200)/cpu/reverse/Flux/(200, 128)`	`273996` ns	`276438` ns	`0.99`
`Dense(200 => 200)/cpu/reverse/SimpleChains/(200, 128)`	`405704` ns	`414357` ns	`0.98`
`Dense(200 => 200)/cpu/reverse/Enzyme/(200, 128)`	`83307` ns	`83792.5` ns	`0.99`
`Dense(200 => 200)/cpu/forward/NamedTuple/(200, 128)`	`80972` ns	`81814` ns	`0.99`
`Dense(200 => 200)/cpu/forward/ComponentArray/(200, 128)`	`81163` ns	`82705` ns	`0.98`
`Dense(200 => 200)/cpu/forward/Flux/(200, 128)`	`86543` ns	`87064` ns	`0.99`
`Dense(200 => 200)/cpu/forward/SimpleChains/(200, 128)`	`104356` ns	`104616` ns	`1.00`
`Conv((3, 3), 16 => 16)/cpu/reverse/ReverseDiff (compiled)/(64, 64, 16, 128)`	`194512141` ns	`192507909.5` ns	`1.01`
`Conv((3, 3), 16 => 16)/cpu/reverse/Zygote/(64, 64, 16, 128)`	`327565451.5` ns	`323234152` ns	`1.01`
`Conv((3, 3), 16 => 16)/cpu/reverse/Tracker/(64, 64, 16, 128)`	`440901510` ns	`413620121` ns	`1.07`
`Conv((3, 3), 16 => 16)/cpu/reverse/ReverseDiff/(64, 64, 16, 128)`	`494489012` ns	`460075216.5` ns	`1.07`
`Conv((3, 3), 16 => 16)/cpu/reverse/Flux/(64, 64, 16, 128)`	`413745131` ns	`388836774` ns	`1.06`
`Conv((3, 3), 16 => 16)/cpu/reverse/SimpleChains/(64, 64, 16, 128)`	`339166533` ns	`355864946` ns	`0.95`
`Conv((3, 3), 16 => 16)/cpu/reverse/Enzyme/(64, 64, 16, 128)`	`44198027` ns	`44389027` ns	`1.00`
`Conv((3, 3), 16 => 16)/cpu/forward/NamedTuple/(64, 64, 16, 128)`	`44263179` ns	`44440540.5` ns	`1.00`
`Conv((3, 3), 16 => 16)/cpu/forward/ComponentArray/(64, 64, 16, 128)`	`43938899` ns	`44195892` ns	`0.99`
`Conv((3, 3), 16 => 16)/cpu/forward/Flux/(64, 64, 16, 128)`	`70454269.5` ns	`60815754` ns	`1.16`
`Conv((3, 3), 16 => 16)/cpu/forward/SimpleChains/(64, 64, 16, 128)`	`27809352` ns	`28722749` ns	`0.97`
`Dense(2000 => 2000)/cpu/reverse/ReverseDiff (compiled)/(2000, 128)`	`19004038` ns	`19205844.5` ns	`0.99`
`Dense(2000 => 2000)/cpu/reverse/Zygote/(2000, 128)`	`19574037` ns	`19859609` ns	`0.99`
`Dense(2000 => 2000)/cpu/reverse/Tracker/(2000, 128)`	`23380541` ns	`23705159.5` ns	`0.99`
`Dense(2000 => 2000)/cpu/reverse/ReverseDiff/(2000, 128)`	`24167625` ns	`24353181.5` ns	`0.99`
`Dense(2000 => 2000)/cpu/reverse/Flux/(2000, 128)`	`19682363` ns	`19858547` ns	`0.99`
`Dense(2000 => 2000)/cpu/reverse/Enzyme/(2000, 128)`	`6557212` ns	`6604286` ns	`0.99`
`Dense(2000 => 2000)/cpu/forward/NamedTuple/(2000, 128)`	`6504547.5` ns	`6592334.5` ns	`0.99`
`Dense(2000 => 2000)/cpu/forward/ComponentArray/(2000, 128)`	`6502118` ns	`6529285` ns	`1.00`
`Dense(2000 => 2000)/cpu/forward/Flux/(2000, 128)`	`6510414` ns	`6570212.5` ns	`0.99`