propertynames of CA from type

Project.toml

@@ -1,7 +1,7 @@
 name = "Lux"
 uuid = "b2108857-7c20-44ae-9111-449ecde12c47"
 authors = ["Avik Pal <avikpal@mit.edu> and contributors"]
-version = "0.4.52"
+version = "0.4.53"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
@@ -43,7 +43,7 @@ LuxZygoteExt = "Zygote"
 Adapt = "3"
 ChainRulesCore = "1"
 ComponentArrays = "0.13"
-FillArrays = "0.13"
+FillArrays = "0.13, 1"
 Flux = "0.13"
 Functors = "0.2, 0.3, 0.4"
 LuxCUDA = "0.1"

ext/LuxComponentArraysExt.jl

@@ -6,8 +6,11 @@ using Functors, Lux, Optimisers
 import TruncatedStacktraces: @truncate_stacktrace
 import ChainRulesCore as CRC
 
-@inline function Lux._getproperty(x::ComponentArray, ::Val{prop}) where {prop}
-    return prop in propertynames(x) ? getproperty(x, prop) : nothing
+@generated function Lux._getproperty(x::ComponentArray{T, N, A, Tuple{Ax}},
+                                     ::Val{v}) where {v, T, N, A,
+                                                      Ax <: ComponentArrays.AbstractAxis}
+    names = propertynames(ComponentArrays.indexmap(Ax))
+    return v ∈ names ? :(x.$v) : :(nothing)
 end
 
 function Functors.functor(::Type{<:ComponentArray}, c)

src/utils.jl

@@ -213,11 +213,12 @@ end
 
 # If doesn't have a property, return nothing
 @generated function _getproperty(x::NamedTuple{names}, ::Val{v}) where {names, v}
-    if v in names
-        return :(x.$v)
-    else
-        return :(nothing)
-    end
+    return v ∈ names ? :(x.$v) : :(nothing)
+end
+
+## Slow-fallback
+@inline function _getproperty(x, ::Val{v}) where {v}
+    return v ∈ propertynames(x) ? getproperty(x, v) : nothing
 end
 
 @inline function _eachslice(x::AbstractArray, ::Val{dims}) where {dims}