Merge pull request #547 from LuxDL/ap/docs_gadjoint

Minor fixes

Project.toml

@@ -1,7 +1,7 @@
 name = "Lux"
 uuid = "b2108857-7c20-44ae-9111-449ecde12c47"
 authors = ["Avik Pal <[email protected]> and contributors"]
-version = "0.5.23"
+version = "0.5.24"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

docs/tutorials.jl

@@ -87,5 +87,5 @@ try
     end
 catch e
     rmprocs(workers()...)
-    @error e
+    rethrow(e)
 end

ext/LuxComponentArraysExt.jl

@@ -1,6 +1,6 @@
 module LuxComponentArraysExt
 
-using ComponentArrays, Functors, Lux, Optimisers
+using Adapt, ComponentArrays, Functors, Lux, Optimisers
 import TruncatedStacktraces: @truncate_stacktrace
 import ChainRulesCore as CRC
 

ext/LuxSimpleChainsExt.jl

@@ -7,7 +7,15 @@ import Lux: SimpleChainsModelConversionError, __to_simplechains_adaptor,
 import Optimisers
 
 function __fix_input_dims_simplechain(layers::Vector, input_dims)
-    return SimpleChains.SimpleChain(input_dims, layers...)
+    L = Tuple(layers)
+    return SimpleChains.SimpleChain{typeof(input_dims), typeof(L)}(input_dims, L)
+end
+
+function __fix_input_dims_simplechain(layers, input_dims)
+    @warn "The model provided is not a `Chain`. Trying to wrap it into a `Chain` but this \
+           might fail. Please consider using `Chain` directly (potentially with \
+           `disable_optimizations = true`)."
+    return __fix_input_dims_simplechain([layers], input_dims)
 end
 
 __equivalent_simplechains_fn(::typeof(Lux.relu)) = SimpleChains.relu
@@ -75,18 +83,4 @@ function NNlib.logsoftmax!(y::SimpleChains.StrideArray{T1, 2},
     return y
 end
 
-# Nicer Interactions with Optimisers.jl
-# function Optimisers._setup(opt::Optimisers.AbstractRule,
-#         ps::Union{SimpleChains.StrideArray, SimpleChains.PtrArray}; cache)
-#     ℓ = Leaf(rule, init(rule, x))
-#     if isbits(x)
-#       cache[nothing] = nothing  # just to disable the warning
-#       ℓ
-#     else
-#       cache[x] = ℓ
-#     end
-#     error(1)
-#     return Optimisers.setup(opt, ps .- ps)
-# end
-
 end

src/utils.jl

@@ -189,6 +189,9 @@ end
 struct LuxEltypeAdaptor{T} end
 
 (l::LuxEltypeAdaptor)(x) = fmap(adapt(l), x)
+function (l::LuxEltypeAdaptor)(x::AbstractArray{T}) where {T}
+    return isbitstype(T) ? adapt(l, x) : map(adapt(l), x)
+end
 
 function Adapt.adapt_storage(
         ::LuxEltypeAdaptor{T}, x::AbstractArray{<:AbstractFloat}) where {T <: AbstractFloat}
@@ -212,6 +215,26 @@ for (fname, ftype) in zip((:f16, :f32, :f64), (Float16, Float32, Float64))
     end
 end
 
+# Common incorrect usage
+for f in (f16, f32, f64)
+    warn_msg = "$(f) is not meant to be broadcasted like `$(f).(x)` or `x .|> $(f)`, and \
+                this might give unexpected results and could lead to crashes. Directly use \
+                `$(f)` as `$(f)(x)` or `x |> $(f)` instead."
+    @eval begin
+        function Base.Broadcast.broadcasted(::typeof($(f)), arg1)
+            @warn $(warn_msg)
+            arg1′ = Broadcast.broadcastable(arg1)
+            return Broadcast.broadcasted(Broadcast.combine_styles(arg1′), $(f), arg1′)
+        end
+
+        function Base.Broadcast.broadcasted(::typeof(|>), arg1, ::typeof($(f)))
+            @warn $(warn_msg)
+            arg1′ = Broadcast.broadcastable(arg1)
+            return Broadcast.broadcasted(Broadcast.combine_styles(arg1′), $(f), arg1′)
+        end
+    end
+end
+
 # Used in freezing
 ## Extend for custom types
 @inline _pairs(x) = pairs(x)

test/transform/simple_chains_tests.jl

@@ -47,4 +47,21 @@
     gs = Zygote.gradient((x, p) -> sum(first(simple_chains_model(x, p, st))), x, ps)
     @test size(gs[1]) == size(x)
     @test length(gs[2].params) == length(ps.params)
+
+    @testset "Single Layer Conversion: LuxDL/Lux.jl#545" begin
+        lux_model = Dense(10 => 5)
+
+        adaptor = ToSimpleChainsAdaptor((static(10),))
+
+        simple_chains_model = @test_warn "The model provided is not a `Chain`. Trying to wrap it into a `Chain` but this might fail. Please consider using `Chain` directly (potentially with `disable_optimizations = true`)." adaptor(lux_model)
+
+        ps, st = Lux.setup(Random.default_rng(), simple_chains_model)
+
+        x = randn(Float32, 10, 3)
+        @test size(first(simple_chains_model(x, ps, st))) == (5, 3)
+
+        gs = Zygote.gradient((x, p) -> sum(first(simple_chains_model(x, p, st))), x, ps)
+        @test size(gs[1]) == size(x)
+        @test length(gs[2].params) == length(ps.params)
+    end
 end