Shouldn't we assert this regardless of the value of `dims`? I suggest you put this next to the other assert in this function.