Skip to content

Add support for custom adjacency context in masked autoregressive transform #65

Merged
francois-rozet merged 4 commits into
probabilists:masterfrom
aalmodovares:context-adj-matrix
Apr 13, 2025
Merged

Add support for custom adjacency context in masked autoregressive transform #65
francois-rozet merged 4 commits into
probabilists:masterfrom
aalmodovares:context-adj-matrix

Conversation

@aalmodovares
Copy link
Copy Markdown
Contributor

This PR implements #64, adding support for user-defined structural constraints in the context features of MaskedAutoregressiveTransform.

Allows specifying an adjacency matrix of shape (features, features + context).

The right-most columns correspond to context-related constraints.

Adds assertions to verify shape validity.

Includes a test to ensure correct behavior.

No additional assumptions are made about the structure of the context adjacency beyond its shape.

Feel free to suggest any changes or additional constraints to enforce.

Copy link
Copy Markdown
Member

@francois-rozet francois-rozet left a comment

Choose a reason for hiding this comment

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

Hi @aalmodovares, I reviewed the proposed changes. These look very appropriate, and the tests as well. Good job! I left a few comments, mainly regarding variable names and a bit of formatting. Tell me if you agree with the suggested changes. I will merge right after. Thanks a lot!

Comment thread zuko/flows/autoregressive.py Outdated
Comment thread tests/test_flows.py Outdated
Comment thread tests/test_flows.py Outdated
Comment thread tests/test_flows.py Outdated
Comment thread tests/test_flows.py Outdated
Comment thread zuko/flows/autoregressive.py Outdated
Comment thread zuko/flows/autoregressive.py Outdated
Comment thread zuko/flows/autoregressive.py Outdated
Comment thread zuko/flows/autoregressive.py Outdated
@aalmodovares
Copy link
Copy Markdown
Contributor Author

Hi @francois-rozet. Thank you for accepting my PR!

I agree with all your suggestions, I have implemented them already and checked everything works. All the suggestions are solved.

Thank you also for the brevity. Feel free to reach me again if something can be improved! :)

Comment thread zuko/flows/autoregressive.py Outdated
Comment thread tests/test_flows.py
@francois-rozet francois-rozet changed the title ✨ Add support for customized adjacency context in masked autoregressive transform Add support for custom adjacency context in masked autoregressive transform Apr 13, 2025
@francois-rozet francois-rozet merged commit 313481e into probabilists:master Apr 13, 2025
@francois-rozet
Copy link
Copy Markdown
Member

The PR has been merged 🥳 Very good job on your first contribution to Zuko!

@aalmodovares
Copy link
Copy Markdown
Contributor Author

Thank you!! I would be glad to collaborate again!

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