Code example in `vmap` notebook is incorrect.

tl;dr version to copy/paste into a corrected version of a notebook:

```python
def vmapped_func(array):
    result = []
    for element in array:
        result.append(func(element))
    result = np.stack(result)
    return result
```