Improve jit function docstring