Remove model_ref in token_graph_builder_model#342
Merged
boomanaiden154 merged 5 commits intomainfrom May 19, 2025
Merged
Conversation
Created using spr 1.3.4 [skip ci]
Created using spr 1.3.4
boomanaiden154
added a commit
to boomanaiden154/gematria
that referenced
this pull request
May 17, 2025
With the port to TF2 I added a hack to get access to instruction annotations and the node mask within the call function of TokenGraphBuilderModelNodeEmbed. This was necessary due to TF2 preferring eager mode so we could not just pass tensor references around. This patch makes everything more canonically TF2 by moving the data around through standard dataflow (with some slight complexity to ensure things are getting passed around to the right places) rather than passing references around. Either one should theoretically work with tf.function annotations depending upon scope, but this will definitely work. Pull Request: google#342
ondrasej
approved these changes
May 19, 2025
Collaborator
ondrasej
left a comment
ondrasej
left a comment
There was a problem hiding this comment.
Cool!
I think we could simplify the models quite a bit by removing the layers abstraction, and writing this directly in the model. But that is for some other time.
Created using spr 1.3.4 [skip ci]
boomanaiden154
changed the base branch from
users/boomanaiden154/main.remove-model_ref-in-token_graph_builder_model
to
main
boomanaiden154
deleted the
users/boomanaiden154/remove-model_ref-in-token_graph_builder_model
branch
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
With the port to TF2 I added a hack to get access to instruction
annotations and the node mask within the call function of
TokenGraphBuilderModelNodeEmbed. This was necessary due to TF2
preferring eager mode so we could not just pass tensor references
around.
This patch makes everything more canonically TF2 by moving the data
around through standard dataflow (with some slight complexity to ensure
things are getting passed around to the right places) rather than
passing references around. Either one should theoretically work with
tf.function annotations depending upon scope, but this will definitely
work.