Skip to content

mvn tri distribution#11

Merged
lockwo merged 2 commits intolockwo:mainfrom
mayalenE:mvn_tri
Jun 10, 2024
Merged

mvn tri distribution#11
lockwo merged 2 commits intolockwo:mainfrom
mayalenE:mvn_tri

Conversation

@mayalenE
Copy link
Copy Markdown
Contributor

@mayalenE mayalenE commented Jun 9, 2024

Hi!

Thank you for initiating this project, exactly what I was looking for 😀
In this PR i've intended to add the mvn tri distribution (and i'm planning to add the one hot categorical soon)

I'm not familiar with distrax so i'm not sure if I did everything correctly, in particular in dropping the batch dimension in TriangularLinear bijector.

However I've ran some basic tests where I check whether logprobs return same value than distrax's logprobs and seems to work well, but might need to do more thorough checking.

Copy link
Copy Markdown
Owner

@lockwo lockwo left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks so much for this PR! I think it's a really great start.

Comment thread distreqx/bijectors/triangular_linear.py Outdated
Comment thread distreqx/bijectors/triangular_linear.py Outdated
Comment thread distreqx/bijectors/triangular_linear.py
Comment thread distreqx/bijectors/triangular_linear.py
Comment thread distreqx/distributions/mvn_tri.py Outdated
Comment thread tests/mvn_tri_test.py
Comment thread tests/mvn_tri_test.py
Comment thread distreqx/bijectors/__init__.py
@mayalenE
Copy link
Copy Markdown
Contributor Author

Hi @lockwo thank you for the detailed reply :)

I've added all the requested changes, except the vmap which I've left to vectorize for the moment, let me know if there is anything else I can do.

@lockwo lockwo self-requested a review June 10, 2024 20:22
Copy link
Copy Markdown
Owner

@lockwo lockwo left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM, thanks for the quick turnaround

@lockwo lockwo merged commit 7bb389b into lockwo:main Jun 10, 2024
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants