Add a few sanity checks

src/functor.jl

@@ -54,6 +54,8 @@ function params(m...)
 end
 
 function loadparams!(m, xs)
+  length(params(m)) == length(xs) ||
+    error("Expected $(length(params(m))) parameter arrays, but got $(length(xs))")
   for (p, x) in zip(params(m), xs)
     size(p) == size(x) ||
       error("Expected param size $(size(p)), got $(size(x))")

src/layers/basic.jl

@@ -35,7 +35,7 @@ functor(::Type{<:Chain}, c) = c.layers, ls -> Chain(ls...)
 applychain(::Tuple{}, x) = x
 applychain(fs::Tuple, x) = applychain(tail(fs), first(fs)(x))
 
-(c::Chain)(x) = applychain(c.layers, x)
+(c::Chain)(x) = applychain(c.layers, x) |> _chain_typecheck
 
 Base.getindex(c::Chain, i::AbstractArray) = Chain(c.layers[i]...)
 
@@ -47,6 +47,15 @@ function Base.show(io::IO, c::Chain)
   print(io, ")")
 end
 
+_chain_typecheck(y) = @show y
+_chain_typewarn(y, ∇y) = @warn "Chain receives gradient of mismatched element type, this may be slow!" typeof(y) typeof(∇y) maxlog=1
+
+Zygote.@adjoint _chain_typecheck(y) = y, ∇y -> begin
+  eltype(y) == eltype(∇y) || _chain_typewarn(y, ∇y)
+  (∇y,)
+end
+Zygote.@nograd _chain_typewarn
+
 """
     outdims(c::Chain, isize)
 

src/utils.jl

@@ -391,6 +391,8 @@ function _restructure(m, xs)
     i += length(x)
     return x
   end
+  i == length(xs) || error("Expected a vector of $i parameters, got $(length(xs))")
+  m
 end
 
 @adjoint function _restructure(m, xs)

test/layers/basic.jl

@@ -121,4 +121,18 @@ import Flux: activations
     m = Maxout(() -> Conv((3, 3), 3 => 16), 2)
     @test Flux.outdims(m, (10, 10)) == (8, 8)
   end
+
+  @testset "type mismatches" begin
+    m = Chain(Dense(2, 10, x->x + 0.1), Dense(10, 2)) # activation creates Float64
+    x = rand(Float32, 2); y=rand(Float32, 2)
+    @test_logs (:warn, # warning from 2nd layer, forward pass
+      "Element types don't match for layer Dense(10, 2), this will be slow."
+      ) gradient(() -> sum(m(x) .- y), params(m))
+
+    m = Chain(Dense(2, 2))
+    x = rand(Float32, 2); y = rand(2) # loss function contains Float64
+    @test_logs (:warn,
+      "Chain receives gradient of mismatched element type, this may be slow!"
+      ) gradient(() -> sum(abs2, m(x) .- y), params(m))
+  end
 end

test/utils.jl

@@ -272,4 +272,15 @@ end
       testdense(re(p), bt)
     end
   end
-end 
+end 
+@testset "size checks" begin
+  m = Chain(Dense(5,4), Dense(4,3))
+  p, re = destructure(m)
+  @test_throws BoundsError re(p[1:end-1])
+  @test_throws ErrorException re(vcat(p,1))
+
+  @test_throws ErrorException loadparams!(Chain(Dense(5,4), Dense(4,3), Dense(3,2)), params(m))
+  @test_throws ErrorException loadparams!(Chain(Dense(5,4), Dense(4,3, bias=false)), params(m))
+  @test_throws ErrorException loadparams!(Chain(Dense(5,4, bias=false), Dense(4,3)), params(m))
+  @test_throws ErrorException loadparams!(Chain(Dense(3,4), Dense(4,5)), params(m))
+end