While the compiled function runs much faster due to type instability, for bigger tuples this leads to a blowup of compilation time for Zygote when differentiating through the unflatten function defined recursively (unsurprisingly perhaps). We should maybe add an rrule.
While the compiled function runs much faster due to type instability, for bigger tuples this leads to a blowup of compilation time for Zygote when differentiating through the
unflattenfunction defined recursively (unsurprisingly perhaps). We should maybe add anrrule.