WARNING: Compilation is currently very experimental and example here don't work out of the box with functorch