Add a leaf_class parameter to flatten (#142)

* add a leaf_class parameter to flatten

* add some docs to flatten for leaf_class