Jittable transforms #1575
Jittable transforms #1575
Conversation
There was a problem hiding this comment.
Thanks so much @pierreglaser! I just have a small nit. Super nice PR!!!
|
Supporting all distributions is complicated I would say. But with jittable support/transform, things might be easier, hopefully. Thanks for working on it!! |
|
Hi @pierreglaser, I would like to make a release this week. Do you want to have this PR in? |
|
Yes, let me finish the PR today and we can merge this. |
|
Ok @fehiepsi, I addressed #1575 (comment) and added a few more tests to make sure that one could |
|
EDIT: it did. |
numpyro/distributions/constraints.py
Outdated
|
|
||
| class _Positive(_GreaterThan, _SingletonConstraint): | ||
| def __eq__(self, other): | ||
| return isinstance(other, _GreaterThan) & jnp.array_equal( |
There was a problem hiding this comment.
Mmh, I think the use of the "bitwise and" operator as a jax-compatible and (that does not try to coerce an abstract boolean to a concrete value) is causing __eq__ to attempt accessing attribute like lower_bound even when other is not of type _GreaterThan. For some reason this happens during the tests, but it could have happened otherwise, so good the tests caught it!
|
OK @fehiepsi, I pushed a fix to the doc errors, that actually revealed some lack of robustness in the updated equality checks for constraints and transform in this PR. I added tests so hopefully that won't be a problem in the future :-) |
* [WIP] jittable transforms * add licence to new test file * turn BijectorConstraint into pytree * test flattening/unflattening of parametrized constraints * cosmetic edits * fix typo * implement tree_flatten/unflatten for transforms * attempt to avoid confusing black * add (un)flattening meths for BijectorTransform * fixup! implement tree_flatten/unflatten for transforms * test vmapping over transforms/constraints * Make constraints `__eq__` checks robust to arbitrary inputs * make transforms equality check robust to arbitrary inputs * test constraints and transforms equality checks
Standalone PR that will vastly ease the flattening/unflattening logic of
TransformedDistributions that will eventually get implemented in #1529 or one of its follow-ups.cc @fehiepsi.