Fixing order of jit and hk.transform

[tomatillos]

Current demos are throwing errors, changed to jitting the .apply method and everything works again.