Restructure on Dense no longer plays nicely with alternative types

[ChrisRackauckas]

MWE:

using ReverseDiff, Flux
ann = Chain(Dense(2,10,tanh), Dense(10,1))
p1,re = Flux.destructure(ann)
function f(p)
    re(p[1:41])([1.0,2.0])
end

ReverseDiff.gradient(f,p3)

MethodError: ReverseDiff.TrackedReal{Float32, Float32, ReverseDiff.TrackedArray{Float32, Float32, 2, Matrix{Float32}, Matrix{Float32}}}(::ForwardDiff.Dual{ForwardDiff.Tag{ReverseDiff.var"#105#107"{DataType, Tuple{}, Val{(1,)}}, Float32}, Float32, 1}) is ambiguous. Candidates:
  (T::Type{var"#s31"} where var"#s31"<:Real)(x::ForwardDiff.Dual) in Tracker at C:\Users\accou\.julia\packages\Tracker\YNNTM\src\lib\real.jl:111
  ReverseDiff.TrackedReal{V, D, O}(value) where {V, D, O} in ReverseDiff at C:\Users\accou\.julia\packages\ReverseDiff\iHmB4\src\tracked.jl:56
Possible fix, define
  ReverseDiff.TrackedReal{V, D, O}(::ForwardDiff.Dual) where {V, D, O}
macro expansion at broadcast.jl:126 [inlined]
splatcall at broadcast.jl:111 [inlined]
(::ReverseDiff.var"#105#107"{DataType, Tuple{}, Val{(1,)}})(s::StaticArrays.SVector{1, ForwardDiff.Dual{ForwardDiff.Tag{ReverseDiff.var"#105#107"{DataType, Tuple{}, Val{(1,)}}, Float32}, Float32, 1}}) at broadcast.jl:150
static_dual_eval at apiutils.jl:32 [inlined]
vector_mode_gradient! at gradient.jl:125 [inlined]
gradient! at gradient.jl:48 [inlined]
(::ReverseDiff.var"#df#106"{DataType, DiffResults.ImmutableDiffResult{1, Float32, Tuple{StaticArrays.SVector{1, Float32}}}, Tuple{}, Val{(1,)}})(x::Float32) at broadcast.jl:148
_broadcast_getindex_evalf at broadcast.jl:648 [inlined]
_broadcast_getindex at broadcast.jl:621 [inlined]
getindex at broadcast.jl:575 [inlined]
copy at broadcast.jl:922 [inlined]
materialize at broadcast.jl:883 [inlined]
broadcast(f::ReverseDiff.var"#df#106"{DataType, DiffResults.ImmutableDiffResult{1, Float32, Tuple{StaticArrays.SVector{1, Float32}}}, Tuple{}, Val{(1,)}}, As::Vector{Float32}) at broadcast.jl:821
∇broadcast at broadcast.jl:154 [inlined]
copy(_bc::Base.Broadcast.Broadcasted{ReverseDiff.TrackedStyle, Tuple{Base.OneTo{Int64}}, Type{ReverseDiff.TrackedReal{Float32, Float32, ReverseDiff.TrackedArray{Float32, Float32, 2, Matrix{Float32}, Matrix{Float32}}}}, Tuple{ReverseDiff.TrackedArray{Float32, Float32, 1, Vector{Float32}, Vector{Float32}}}}) at broadcast.jl:94
materialize(bc::Base.Broadcast.Broadcasted{ReverseDiff.TrackedStyle, Nothing, Type{ReverseDiff.TrackedReal{Float32, Float32, ReverseDiff.TrackedArray{Float32, Float32, 2, Matrix{Float32}, Matrix{Float32}}}}, Tuple{ReverseDiff.TrackedArray{Float32, Float32, 1, Vector{Float32}, Vector{Float32}}}}) at broadcast.jl:883
broadcast at broadcast.jl:821 [inlined]
create_bias(weights::ReverseDiff.TrackedArray{Float32, Float32, 2, Matrix{Float32}, Matrix{Float32}}, bias::ReverseDiff.TrackedArray{Float32, Float32, 1, Vector{Float32}, Vector{Float32}}, dims::Int64) at utils.jl:397
Dense(W::ReverseDiff.TrackedArray{Float32, Float32, 2, Matrix{Float32}, Matrix{Float32}}, bias::ReverseDiff.TrackedArray{Float32, Float32, 1, Vector{Float32}, Vector{Float32}}, σ::typeof(tanh)) at basic.jl:117
(::Flux.var"#126#127")(y::NamedTuple{(:weight, :bias, :σ), Tuple{ReverseDiff.TrackedArray{Float32, Float32, 2, Matrix{Float32}, Matrix{Float32}}, ReverseDiff.TrackedArray{Float32, Float32, 1, Vector{Float32}, Vector{Float32}}, typeof(tanh)}}) at functor.jl:18
fmap1(f::Function, x::Dense{typeof(tanh), Matrix{Float32}, Vector{Float32}}) at functor.jl:49
fmap(f::Flux.var"#54#55"{ReverseDiff.TrackedArray{Float32, Float32, 1, Vector{Float32}, Vector{Float32}}}, x::Dense{typeof(tanh), Matrix{Float32}, Vector{Float32}}; exclude::typeof(Functors.isleaf), cache::IdDict{Any, Any}) at functor.jl:56

The stack trace is a bit off and gets cut off, but the issue comes from:

function create_bias(weights::AbstractArray, bias::AbstractArray, dims::Integer...)
  size(bias) == dims || throw(DimensionMismatch("expected bias of size $(dims), got size $(size(bias))"))
  if eltype(bias) == eltype(weights)
    return bias
  else
    @warn "converting bias to match element type of weights" typeof(weights) typeof(bias) maxlog=3 _id=hash(dims)
    return broadcast(eltype(weights), bias)
  end
end

Generally that's not the most correct way to convert types. A better way to handle this is:

function Flux.create_bias(weights::AbstractArray, bias::AbstractArray, dims::Integer...)
  size(bias) == dims || throw(DimensionMismatch("expected bias of size $(dims), got size $(size(bias))"))
  if eltype(bias) == eltype(weights)
    return bias
  else
    @warn "converting bias to match element type of weights" typeof(weights) typeof(bias) maxlog=3 _id=hash(dims)
    return convert.((eltype(weights),),bias)
  end
end

However, that's only a partial solution since it shouldn't be throwing a warning on a legitimate use case, so there's a bigger question as to why it ends up there in the first place (which is cut off from the stack trace).

[DhairyaLGandhi]

Yeah I agree. We can't have restrictions on the kinds of types we have to work with. More importantly, its not correct to convert Tracked* style of types. We need them to track the ops. In other cases, its required to actually convert the parameters to dispatch correctly.

This is generally too restrictive and should be reverted to match original behaviour.

[ChrisRackauckas]

For reference:

eltype(weights) = ReverseDiff.TrackedReal{Float32, Float32, ReverseDiff.TrackedArray{Float32, Float32, 2, Matrix{Float32}, Matrix{Float32}}}
eltype(bias) = ReverseDiff.TrackedReal{Float32, Float32, ReverseDiff.TrackedArray{Float32, Float32, 1, Vector{Float32}, Vector{Float32}}}

The only difference is the dimensionality of the array that they come from, which is because of how Flux's fmap is splitting it. It really should just be creating Dense(W,b,olddense.activation) kind of deal.

[ChrisRackauckas]

Is there a technical reason why restructure isn't just using a view?

[DhairyaLGandhi]

I'm not sure. I don't think there's a major reason why that wouldn't be okay. It would be much better to avoid the errant code path entirely.

[darsnack]

Looks like that restriction was put in place because of #1421 (ref: #1440 (comment)). There's some further discussion about the restriction in that thread.

[darsnack]

I disagree with the reasoning in that thread. If you want the layer to promote then specify types that will result in promotion. If you want to avoid promotion, then specify homogenous types. The restriction should probably be relaxed to a warning on construction (lot's of people get hit by the promotion mismatch).

[ChrisRackauckas]

Seems out of place. That's for conv, presumably because the wrapper doesn't handle different types? But linear algebra doesn't have that same issue, so throwing Dense under the same bus doesn't seem right.

[darsnack]

Yeah if conv has a harder restriction, then we should address it specifically in the construction or forward pass of Conv.

[darsnack]

The error should be moved to here. Are we fine with a warning though (in the place of forcing the conversion)? It would help people who run into this problem without being too restrictive. If we're happy with that, I can submit a PR later this evening (if someone else doesn't get to it first).

[DhairyaLGandhi]

I think we could rethink the create_bias kinds of structures. Promotion isn't a bad thing in and of itself, and in case we want to restrict mismatched types, its much clearer to have that error.

[DhairyaLGandhi]

There is no limitation for Conv to use homogenous types either. We have a fallback method in NNlib that already emits the mismatch types warning too, with an informative blurb about which types don't match.

[darsnack]

It seems like create_bias is just a utility to handle bias = true/bias = false cases (as opposed to replicating the code for every constructor that needs to handle a bias argument. But I agree with you that the way it is written seems to overstep what it should do.

Promotion isn't a bad thing in and of itself, and in case we want to restrict mismatched types, its much clearer to have that error.
We have a fallback method in NNlib that already emits the mismatch types warning too, with an informative blurb about which types don't match.

Okay so sounds like we just drop the eltype checks all together.

[ChrisRackauckas]

Yeah, it seems that by design it can only:

1. Throw an error on things which already throw the same error
2. Throw a warning on cases which actually work

which IMO is kind of odd.

[darsnack]

Yeah it seems like the homogenous type restriction in #1421 is a misdirection anyways. The struct nor NNlib.conv had such a restriction. It looks like it was probably some badly specified types in a user-facing constructor. Fortunately, the recent refactor PRs for the constructors seems to have fixed that.

Yeah, throwing a warning for the TrackedReal mismatch would be really weird.