fix: StaticArray support regression

Project.toml

@@ -1,7 +1,7 @@
 name = "LuxLib"
 uuid = "82251201-b29d-42c6-8e01-566dec8acb11"
 authors = ["Avik Pal <[email protected]> and contributors"]
-version = "0.3.32"
+version = "0.3.33"
 
 [deps]
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
@@ -18,6 +18,7 @@ Markdown = "d6f4376e-aef5-505a-96c1-9c027394607a"
 NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
+StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 UnrolledUtilities = "0fe1646c-419e-43be-ac14-22321958931b"
 
@@ -62,6 +63,8 @@ ReTestItems = "1.23.1"
 Reexport = "1"
 ReverseDiff = "1.15"
 StableRNGs = "1"
+StaticArrays = "1.9"
+StaticArraysCore = "1.4.3"
 Statistics = "1.10"
 Test = "1.10"
 Tracker = "0.2.34"
@@ -82,9 +85,10 @@ Preferences = "21216c6a-2e73-6563-6e65-726566657250"
 ReTestItems = "817f1d60-ba6b-4fd5-9520-3cf149f6a823"
 ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
+StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["Aqua", "ComponentArrays", "Enzyme", "ExplicitImports", "LuxTestUtils", "Pkg", "Preferences", "ReTestItems", "ReverseDiff", "StableRNGs", "Test", "Tracker", "Zygote"]
+test = ["Aqua", "ComponentArrays", "Enzyme", "ExplicitImports", "LuxTestUtils", "Pkg", "Preferences", "ReTestItems", "ReverseDiff", "StableRNGs", "StaticArrays", "Test", "Tracker", "Zygote"]

src/LuxLib.jl

@@ -15,6 +15,7 @@ using Markdown: @doc_str
 using NNlib: NNlib, ConvDims, conv, conv!, relu, gelu, σ, ∇conv_data, ∇conv_filter
 using Random: Random, AbstractRNG, rand!
 using Reexport: @reexport
+using StaticArraysCore: StaticArraysCore, StaticVector
 using Statistics: Statistics, mean, var
 using UnrolledUtilities: unrolled_any, unrolled_all, unrolled_filter
 

src/impl/bias_activation.jl

@@ -1,8 +1,13 @@
 __reshape_bias_into_xdims(::AbstractArray, ::Nothing) = nothing
 __reshape_bias_into_xdims(::AbstractVector, bias::AbstractVector) = bias
-function __reshape_bias_into_xdims(
-        ::AbstractArray{<:Number, N}, bias::AbstractVector) where {N}
-    return reshape(bias, ntuple(i -> ifelse(i == N - 1, length(bias), 1), N))
+__reshape_bias_into_xdims(::AbstractVector, bias::StaticVector) = bias
+function __reshape_bias_into_xdims(x::AbstractArray, bias::AbstractVector)
+    return reshape(bias, ntuple(i -> ifelse(i == ndims(x) - 1, length(bias), 1), ndims(x)))
+end
+function __reshape_bias_into_xdims(x::AbstractArray, bias::StaticVector)
+    return StaticArraysCore.SArray{
+        Tuple{ntuple(i -> ifelse(i == ndims(x) - 1, length(bias), 1), ndims(x))...},
+        eltype(bias), ndims(x), length(bias)}(bias.data)
 end
 
 ## Needed for type stability

test/common_ops/dense_tests.jl

@@ -68,3 +68,13 @@
         end
     end
 end
+
+@testitem "Fused Dense Bias Activation: StaticArrays" tags=[:common_ops] begin
+    using StaticArrays
+
+    x = @SArray rand(2, 4)
+    weight = @SArray rand(3, 2)
+    bias = @SArray rand(3)
+
+    @test @inferred(fused_dense_bias_activation(relu, weight, x, bias)) isa SArray
+end