Add unflatten_like (#200)

* add unflatten_like

* add tests for unflatten_like

* add more checks to unflatten_like

* add docs to motivate unflatten_like

* Add another clarifying sentence to unflatten_like docstring

Co-authored-by: Andreas Klöckner <inform@tiker.net>