danielward27 / flowjax

https://danielward27.github.io/flowjax/
MIT License
82 stars 10 forks source link

Int dummy #160

Closed danielward27 closed 4 months ago

danielward27 commented 4 months ago

Some wrappers use a dummy scalar array to detect vmapped initialization. These were floats, which in some cases is annoying, for example they would be filtered alongside trainable parameters (although their gradient would be zero). I have changed these arrays to integers for now. Fixed a minor bug where if unwrap was called on an object created with a vmapped init, and the AbstractUnwrappable contained no arrays to vectorize over, then the dimensions from the dummy were inferred in the wrong order.