Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[key reuse] refactor & unify key reuse rule registration #20692

Merged
merged 1 commit into from
Apr 12, 2024

Conversation

jakevdp
Copy link
Collaborator

@jakevdp jakevdp commented Apr 10, 2024

This refactoring will pave the way for making key reuse checking part of the core dispatch path. There are two goals:

  1. Move from having two tables of key reuse rules (static and dynamic) to having a single table: this will let us begin easily specifying key reuse signatures at the location where primitives are defined.
  2. Move shared code into single utility functions that can be referenced from the main dispatch path: this will let us move away from the current approach which involves monkey-patching the impl rules.

@jakevdp jakevdp requested a review from froystig April 10, 2024 17:52
@jakevdp jakevdp self-assigned this Apr 10, 2024
@jakevdp jakevdp force-pushed the key-reuse-refactor branch 2 times, most recently from ace6618 to 1f37c1d Compare April 11, 2024 15:12
Copy link
Member

@froystig froystig left a comment

Choose a reason for hiding this comment

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

It seems that we could unify KeyReuseSignature and DynamicKeyReuseSignature by treating all of them as functions of some number of operands in the equation, where KeyReuseSignatures are specifically nullary. Is that worth reflecting in the code (making everything one type, and always a function)?

@google-ml-butler google-ml-butler bot added kokoro:force-run pull ready Ready for copybara import and testing labels Apr 11, 2024
@copybara-service copybara-service bot merged commit 72ab07f into jax-ml:main Apr 12, 2024
13 checks passed
@jakevdp jakevdp deleted the key-reuse-refactor branch April 12, 2024 00:09
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
pull ready Ready for copybara import and testing
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants