The gradients utilizing getproperty takes ages

[pevnak]

In Mill.jl, we use getproperty a lot, and therefore the time to first gradient can be long. Below MWE demonstrates it with the complexity being controlled by n.

using Mill, Flux, Random
Mill.bagcount!(false)

create_array_node() = ArrayNode(Flux.onehotbatch(1:4,1:4))
function create_bagnode_node(create_child)
	BagNode(create_child(), [i:i for i in 1:4])
end

function create_product_node(create_child,n)
	vs = tuple([create_child() for i in 1:n]...)
	ks = tuple([Symbol(randstring()) for i in 1:n]...)
	ProductNode(NamedTuple{ks}(vs))
end

for n in 1:10 
	ds = create_product_node(n) do 
		create_bagnode_node(() -> create_product_node(create_array_node, n))
	end
	model = reflectinmodel(ds)
	ps = Flux.params(model)
	@time gradient(() -> sum(model(ds).data), ps)
end

with compilation times

 39.589115 seconds (73.16 M allocations: 4.323 GiB, 4.09% gc time, 99.84% compilation time)
  6.825684 seconds (5.47 M allocations: 323.489 MiB, 1.25% gc time, 99.05% compilation time)
 12.464049 seconds (7.84 M allocations: 468.132 MiB, 1.06% gc time, 99.08% compilation time)
 15.806830 seconds (10.95 M allocations: 653.072 MiB, 1.26% gc time, 98.74% compilation time)
 22.891749 seconds (14.47 M allocations: 861.895 MiB, 1.27% gc time, 98.50% compilation time)
 33.988718 seconds (18.42 M allocations: 1.070 GiB, 1.24% gc time, 98.55% compilation time)
 49.239070 seconds (22.91 M allocations: 1.330 GiB, 1.07% gc time, 98.41% compilation time)
 70.693351 seconds (27.99 M allocations: 1.624 GiB, 0.95% gc time, 98.29% compilation time)
 99.862504 seconds (33.59 M allocations: 1.947 GiB, 0.82% gc time, 98.18% compilation time)
 215.680960 seconds (38.19 M allocations: 2.202 GiB, 0.45% gc time, 98.50% compilation time)

The environment is

(mwe) pkg> st
     Status `~/Work/Julia/Projects/mwe/Project.toml`
 [587475ba] Flux v0.12.4
 [1d0525e4] Mill v2.5.0
 [e88e6eb3] Zygote v0.6.12

and I have tested it on Julia 1.6.1.

[pevnak]

A relevant PR is #909

[simeonschaub]

#909 has now been merged. Could you check this again with 0.6.27?