Integrate/use flowjax

[lockwo]

cool package I just found https://github.com/danielward27/flowjax/tree/main. Would be interesting to see if they are natively compatible with us or if there is some integration required