danielward27 / flowjax

https://danielward27.github.io/flowjax/
MIT License
101 stars 14 forks source link

Ensure int masks and ranks #110

Closed danielward27 closed 1 year ago

danielward27 commented 1 year ago

Fixes jnp.ones without dtype specified leading to in_ranks being float dtype.