Flux.jl - Best way to do efficient division by a parameter.



The following function evaluation:

function OneDiv(x)
    return  1.0 ./ x
@code_warntype OneDiv(Flux.param([0.1]))

Compiles to:

1 ─ %1 = (Core.tuple)(1.0, x)::Tuple{Float64,TrackedArray{…,Array{Float64,1}}}
│   %2 = %new(Base.Broadcast.Broadcasted{Flux.Tracker.TrackedStyle,Nothing,typeof(/),Tuple{Float64,TrackedArray{…,Array{Float64,1}}}}, /, %1, nothing)::Base.Broadcast.Broadcasted{Flux.Tracker.TrackedStyle,Nothing,typeof(/),Tuple{Float64,TrackedArray{…,Array{Float64,1}}}}
│   %3 = invoke Base.materialize(%2::Base.Broadcast.Broadcasted{Flux.Tracker.TrackedStyle,Nothing,typeof(/),Tuple{Float64,TrackedArray{…,Array{Float64,1}}}})::Any
└──      return %3

Note the Any there. I was wondering how to make this compile to a concrete type and remove the Any here? If I change ./ to /, it fails with: MethodError: no method matching Float64(::Flux.Tracker.TrackedReal{Float64}).

My param in question is actually a scalar, an in that case the code works fine with /, but I changed it to a single-element array because it didn’t seem like train! worked with scalar params. If there is a way for train! to work with scalar params, however, then that’s another option.