Fix unexpected behaviour with km's inputs

[fmeirinhos]

One would still need to fix the docstrings

[LRydin]

@fmeirinhos I think I will close this request. Unfortunately, I cannot see what the improvement is. This is much more explicit than the extant code ─ which isn't bad ─ but this is a very low-level code, so explicit means difficult to understand. I've tried to weave something out of the code, but it is difficult to see where the real improvement is.