add flax pipelines to api doc + doc string examples (#2600)
* add api doc for flax pipeline + doc string examples
* make style
---------
Co-authored-by: yiyixuxu <yixu@yis-macbook-pro.lan>
Co-authored-by: Patrick von Platen <patrick.v.platen@gmail.com>
