Inplace version of Dense

Project.toml

@@ -1,11 +1,12 @@
 name = "Lux"
 uuid = "b2108857-7c20-44ae-9111-449ecde12c47"
 authors = ["Avik Pal <[email protected]> and contributors"]
-version = "0.5.10"
+version = "0.5.11"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
+Bumper = "8ce10254-0962-460f-a3d8-1f77fea1446e"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471"
 Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"

src/Lux.jl

@@ -15,6 +15,9 @@ import ChainRulesCore as CRC
 import TruncatedStacktraces
 import TruncatedStacktraces: @truncate_stacktrace
 
+# Bumper Allocation
+using Bumper
+
 # LuxCore
 @reexport using LuxCore
 import LuxCore: AbstractExplicitLayer, AbstractExplicitContainerLayer, initialparameters,

src/layers/basic.jl

@@ -221,6 +221,16 @@ end
             d.out_dims, size(x)[2:end]...), st)
 end
 
+# Inplace Interface
+@inline function (d::Dense{true})(y::AbstractMatrix, buffer, x::AbstractMatrix, ps,
+        st::NamedTuple)
+    W, b = ps.weight, ps.bias
+    @. y = b
+    mul!(y, W, x, true, true)
+    @. y = d.activation(y)
+    return y, st
+end
+
 """
     Scale(dims, activation=identity; init_weight=ones32, init_bias=zeros32, bias::Bool=true)