* WIP

* cleaner code

* Add tests

* Add docs

* Assert dtype

* put back check

* Address comments
Co-authored-by: Francisco Massa <fvsmassa@gmail.com>