This adds a generated function to efficiently deal with the case, when the keys in NamedTuples of ProductNode and ProductModel do not match. When they match, the old method relying on map still works. If they do not match, there is a generated function which generated efficient code.
Test bench:
using Mill, Flux
a = ArrayNode(randn(Float32, 3, 4))
b = ArrayNode(randn(Float32, 3, 4))
c = ArrayNode(randn(Float32, 3, 4))
x₁ = ProductNode((;a, b, c))
m = reflectinmodel(x₁, d -> Dense(d, 4))
x₂ = ProductNode((;c, b, a))
m(x₁).data ≈ m(x₂).data
The generated function generates code, which is much nicer on zygote. Compile time of a version based on map
This adds a generated function to efficiently deal with the case, when the keys in
NamedTuples
of ProductNode and ProductModel do not match. When they match, the old method relying onmap
still works. If they do not match, there is a generated function which generated efficient code.Test bench:
The generated function generates code, which is much nicer on zygote. Compile time of a version based on
map
Compile time of function that uses generated function
thus we see what has been expected.
Regarding the runtime, both functions seems to be the same (recall
x₁
calls map,x₂
calls generated function)though during the gradient, the generated function is more efficient
threfore I leave up to the committee (@SimonMandlik @racinmat ) if we should try the generated function by default.