Remove type promotion in FFJORD

[prbzrg]

I tried to fix the ReverseDiff issue (#615) with FFJORD by removing the type promotion, but questionably it fixed the ForwardDiff issue (#610).
Now, output type of ffjord is as input type because of #617, I guess. So we don't need to promote the output of logpdf anymore.

Fixes #610

[prbzrg]

Also, this PR changes the error in #615 to this:

regularize=false & monte_carlo=false: Error During Test at /home/runner/work/DiffEqFlux.jl/DiffEqFlux.jl/test/cnf_test.jl:71
  Test threw exception
  Expression: !(isnothing(DiffEqFlux.sciml_train((θ->begin
                    loss(θ; regularize, monte_carlo)
                end), ffjord_mdl.p, ADAM(0.1), adtype; cb, maxiters = 10)))
  DimensionMismatch("arrays could not be broadcast to a common size; got a dimension with lengths 5 and 6")
  Stacktrace:
    [1] _bcs1
      @ ./broadcast.jl:501 [inlined]
    [2] _bcs
      @ ./broadcast.jl:495 [inlined]
    [3] broadcast_shape(::Tuple{Base.OneTo{Int64}}, ::Tuple{Base.OneTo{Int64}})
      @ Base.Broadcast ./broadcast.jl:489
    [4] combine_axes
      @ ./broadcast.jl:484 [inlined]
    [5] _axes
      @ ./broadcast.jl:209 [inlined]
    [6] axes
      @ ./broadcast.jl:207 [inlined]
    [7] copy(bc::Base.Broadcast.Broadcasted{Base.Broadcast.Style{Tuple}, Nothing, typeof(ReverseDiff._add_to_deriv!), Tuple{Tuple{ODEProblem{Matrix{Float64}, Tuple{Float32, Float32}, false, ReverseDiff.TrackedArray{Float32, Float32, 1, Vector{Float32}, Vector{Float32}}, ODEFunction{false, DiffEqFlux.var"#ffjord_#61"{Bool, Bool, FFJORD{Chain{Tuple{Dense{typeof(tanh), Matrix{Float32}, Vector{Float32}}}}, Vector{Float32}, Flux.var"#60#62"{Chain{Tuple{Dense{typeof(tanh), Matrix{Float32}, Vector{Float32}}}}}, Distributions.MvNormal{Float32, PDMats.PDiagMat{Float32, Vector{Float32}}, Vector{Float32}}, Tuple{Float32, Float32}, Tuple{OrdinaryDiffEq.Tsit5}, Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}, Matrix{Float64}}, LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing}, Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, SciMLBase.StandardODEProblem}, InterpolatingAdjoint{0, true, Val{:central}, Bool, Bool}, Matrix{Float64}, ReverseDiff.TrackedArray{Float32, Float32, 1, Vector{Float32}, Vector{Float32}}, OrdinaryDiffEq.Tsit5}, Tuple{ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, ChainRulesCore.NoTangent, Matrix{Float64}, Vector{Float64}, ChainRulesCore.NoTangent}}})
      @ Base.Broadcast ./broadcast.jl:1095
    [8] materialize
      @ ./broadcast.jl:883 [inlined]
    [9] special_reverse_exec!(instruction::ReverseDiff.SpecialInstruction{typeof(DiffEqBase.solve_up), Tuple{ODEProblem{Matrix{Float64}, Tuple{Float32, Float32}, false, ReverseDiff.TrackedArray{Float32, Float32, 1, Vector{Float32}, Vector{Float32}}, ODEFunction{false, DiffEqFlux.var"#ffjord_#61"{Bool, Bool, FFJORD{Chain{Tuple{Dense{typeof(tanh), Matrix{Float32}, Vector{Float32}}}}, Vector{Float32}, Flux.var"#60#62"{Chain{Tuple{Dense{typeof(tanh), Matrix{Float32}, Vector{Float32}}}}}, Distributions.MvNormal{Float32, PDMats.PDiagMat{Float32, Vector{Float32}}, Vector{Float32}}, Tuple{Float32, Float32}, Tuple{OrdinaryDiffEq.Tsit5}, Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}, Matrix{Float64}}, LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing}, Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, SciMLBase.StandardODEProblem}, InterpolatingAdjoint{0, true, Val{:central}, Bool, Bool}, Matrix{Float64}, ReverseDiff.TrackedArray{Float32, Float32, 1, Vector{Float32}, Vector{Float32}}, OrdinaryDiffEq.Tsit5}, ReverseDiff.TrackedArray{Float64, Float64, 3, Array{Float64, 3}, Array{Float64, 3}}, Tuple{DiffEqSensitivity.var"#adjoint_sensitivity_backpass#216"{Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, OrdinaryDiffEq.Tsit5, InterpolatingAdjoint{0, true, Val{:central}, Bool, Bool}, Matrix{Float64}, Vector{Float32}, Tuple{}, Colon, NamedTuple{(), Tuple{}}}, DiffEqBase.var"##solve_up#286#204"{DiffEqBase.var"##solve_up#286#203#205"}, NamedTuple{(), Tuple{}}}})
      @ DiffEqBase ~/.julia/packages/ReverseDiff/E4Tzn/src/macros.jl:218
   [10] reverse_exec!(instruction::ReverseDiff.SpecialInstruction{typeof(DiffEqBase.solve_up), Tuple{ODEProblem{Matrix{Float64}, Tuple{Float32, Float32}, false, ReverseDiff.TrackedArray{Float32, Float32, 1, Vector{Float32}, Vector{Float32}}, ODEFunction{false, DiffEqFlux.var"#ffjord_#61"{Bool, Bool, FFJORD{Chain{Tuple{Dense{typeof(tanh), Matrix{Float32}, Vector{Float32}}}}, Vector{Float32}, Flux.var"#60#62"{Chain{Tuple{Dense{typeof(tanh), Matrix{Float32}, Vector{Float32}}}}}, Distributions.MvNormal{Float32, PDMats.PDiagMat{Float32, Vector{Float32}}, Vector{Float32}}, Tuple{Float32, Float32}, Tuple{OrdinaryDiffEq.Tsit5}, Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}, Matrix{Float64}}, LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing}, Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, SciMLBase.StandardODEProblem}, InterpolatingAdjoint{0, true, Val{:central}, Bool, Bool}, Matrix{Float64}, ReverseDiff.TrackedArray{Float32, Float32, 1, Vector{Float32}, Vector{Float32}}, OrdinaryDiffEq.Tsit5}, ReverseDiff.TrackedArray{Float64, Float64, 3, Array{Float64, 3}, Array{Float64, 3}}, Tuple{DiffEqSensitivity.var"#adjoint_sensitivity_backpass#216"{Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, OrdinaryDiffEq.Tsit5, InterpolatingAdjoint{0, true, Val{:central}, Bool, Bool}, Matrix{Float64}, Vector{Float32}, Tuple{}, Colon, NamedTuple{(), Tuple{}}}, DiffEqBase.var"##solve_up#286#204"{DiffEqBase.var"##solve_up#286#203#205"}, NamedTuple{(), Tuple{}}}})
      @ ReverseDiff ~/.julia/packages/ReverseDiff/E4Tzn/src/tape.jl:93
   [11] reverse_pass!(tape::Vector{ReverseDiff.AbstractInstruction})
      @ ReverseDiff ~/.julia/packages/ReverseDiff/E4Tzn/src/tape.jl:87
   [12] reverse_pass!
      @ ~/.julia/packages/ReverseDiff/E4Tzn/src/api/tape.jl:36 [inlined]
   [13] seeded_reverse_pass!(result::Vector{Float32}, output::ReverseDiff.TrackedReal{Float64, Float64, Nothing}, input::ReverseDiff.TrackedArray{Float32, Float32, 1, Vector{Float32}, Vector{Float32}}, tape::ReverseDiff.GradientTape{GalacticOptim.var"#192#201"{Tuple{}, GalacticOptim.var"#190#199"{OptimizationFunction{true, GalacticOptim.AutoReverseDiff, DiffEqFlux.var"#82#87"{Main.##345.var"#6#28"{Bool, Bool}}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Nothing}}, ReverseDiff.TrackedArray{Float32, Float32, 1, Vector{Float32}, Vector{Float32}}, ReverseDiff.TrackedReal{Float64, Float64, Nothing}})
      @ ReverseDiff ~/.julia/packages/ReverseDiff/E4Tzn/src/api/utils.jl:31
   [14] seeded_reverse_pass!(result::Vector{Float32}, t::ReverseDiff.GradientTape{GalacticOptim.var"#192#201"{Tuple{}, GalacticOptim.var"#190#199"{OptimizationFunction{true, GalacticOptim.AutoReverseDiff, DiffEqFlux.var"#82#87"{Main.##345.var"#6#28"{Bool, Bool}}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Nothing}}, ReverseDiff.TrackedArray{Float32, Float32, 1, Vector{Float32}, Vector{Float32}}, ReverseDiff.TrackedReal{Float64, Float64, Nothing}})
      @ ReverseDiff ~/.julia/packages/ReverseDiff/E4Tzn/src/api/tape.jl:47
   [15] gradient!(result::Vector{Float32}, f::Function, input::Vector{Float32}, cfg::ReverseDiff.GradientConfig{ReverseDiff.TrackedArray{Float32, Float32, 1, Vector{Float32}, Vector{Float32}}})
      @ ReverseDiff ~/.julia/packages/ReverseDiff/E4Tzn/src/api/gradients.jl:42
   [16] (::GalacticOptim.var"#191#200"{GalacticOptim.var"#190#199"{OptimizationFunction{true, GalacticOptim.AutoReverseDiff, DiffEqFlux.var"#82#87"{Main.##345.var"#6#28"{Bool, Bool}}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Nothing}})(::Vector{Float32}, ::Vector{Float32})
      @ GalacticOptim ~/.julia/packages/GalacticOptim/bEh06/src/function/reversediff.jl:9
   [17] macro expansion
      @ ~/.julia/packages/GalacticOptim/bEh06/src/solve/flux.jl:43 [inlined]
   [18] macro expansion
      @ ~/.julia/packages/GalacticOptim/bEh06/src/solve/solve.jl:35 [inlined]
   [19] __solve(prob::OptimizationProblem{false, OptimizationFunction{false, GalacticOptim.AutoReverseDiff, OptimizationFunction{true, GalacticOptim.AutoReverseDiff, DiffEqFlux.var"#82#87"{Main.##345.var"#6#28"{Bool, Bool}}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, GalacticOptim.var"#191#200"{GalacticOptim.var"#190#199"{OptimizationFunction{true, GalacticOptim.AutoReverseDiff, DiffEqFlux.var"#82#87"{Main.##345.var"#6#28"{Bool, Bool}}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Nothing}}, GalacticOptim.var"#193#202"{GalacticOptim.var"#190#199"{OptimizationFunction{true, GalacticOptim.AutoReverseDiff, DiffEqFlux.var"#82#87"{Main.##345.var"#6#28"{Bool, Bool}}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Nothing}}, GalacticOptim.var"#198#207", Nothing, Nothing, Nothing}, Vector{Float32}, SciMLBase.NullParameters, Nothing, Nothing, Nothing, Nothing, Base.Iterators.Pairs{Symbol, typeof(Main.##345.cb), Tuple{Symbol}, NamedTuple{(:cb,), Tuple{typeof(Main.##345.cb)}}}}, opt::ADAM, data::Base.Iterators.Cycle{Tuple{GalacticOptim.NullData}}; maxiters::Int64, cb::Function, progress::Bool, save_best::Bool, kwargs::Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(), Tuple{}}})
      @ GalacticOptim ~/.julia/packages/GalacticOptim/bEh06/src/solve/flux.jl:41
   [20] #solve#476
      @ ~/.julia/packages/SciMLBase/n3U0M/src/solve.jl:3 [inlined]
   [21] sciml_train(::Main.##345.var"#6#28"{Bool, Bool}, ::Vector{Float32}, ::ADAM, ::GalacticOptim.AutoReverseDiff; lower_bounds::Nothing, upper_bounds::Nothing, maxiters::Int64, kwargs::Base.Iterators.Pairs{Symbol, typeof(Main.##345.cb), Tuple{Symbol}, NamedTuple{(:cb,), Tuple{typeof(Main.##345.cb)}}})
      @ DiffEqFlux ~/work/DiffEqFlux.jl/DiffEqFlux.jl/src/train.jl:89
   [22] macro expansion
      @ ~/work/DiffEqFlux.jl/DiffEqFlux.jl/test/cnf_test.jl:71 [inlined]
   [23] macro expansion
      @ /buildworker/worker/package_linux64/build/usr/share/julia/stdlib/v1.6/Test/src/Test.jl:1151 [inlined]
   [24] macro expansion
      @ ~/work/DiffEqFlux.jl/DiffEqFlux.jl/test/cnf_test.jl:68 [inlined]
   [25] macro expansion
      @ /buildworker/worker/package_linux64/build/usr/share/julia/stdlib/v1.6/Test/src/Test.jl:1151 [inlined]
   [26] macro expansion
      @ ~/work/DiffEqFlux.jl/DiffEqFlux.jl/test/cnf_test.jl:65 [inlined]
   [27] macro expansion
      @ /buildworker/worker/package_linux64/build/usr/share/julia/stdlib/v1.6/Test/src/Test.jl:1151 [inlined]
   [28] top-level scope
      @ ~/work/DiffEqFlux.jl/DiffEqFlux.jl/test/cnf_test.jl:22

[prbzrg]

DimensionMismatch("arrays could not be broadcast to a common size; got a dimension with lengths 5 and 6")

The error looks similar to the Tracker issue (#624):

DimensionMismatch("array could not be broadcast to match destination")

[ChrisRackauckas]

Thanks!