Allow stateful function in dataset.map
Removes the "test type" section in Dataset.map which would modify the state of the stateful function. Now, the return type of the map function is inferred after processing the first example.

Fixes #1940 

@lhoestq Not very happy with the usage of `nonlocal`. Would like to hear your opinion on this.