That seems more like the tracing overhead than the python overhead. The original jit proposal would I believe not help at all with that since ML workloads basically do their own jit. The post being discussed however pushes for a general framework and good tracing support so might help ML workloads.