forked from keras-team/keras
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Fix variables binding in JAX export. (keras-team#19399)
In the current JAX export, while the variables are exported as `tf.Variable`s, these `tf.Variable`s are not at all connected to the graph. Instead, all the weights of the variables are inlined in the graph as constants. As a consequence: - The export is almost twice the size it should be - One easily runs into the 2GB limit for the graph size with models with a lot of parameters - Models that use RNGs in inference mode fail to export with a "leak" error. This PR fixes the issues above by creating a stateless function, jax2tf converting it and binding the variables to the function in a `tf.function` wrapper. Also: - Fixed an issue where with the JAX export, the Python signature of the exported function was lost
- Loading branch information
Showing
2 changed files
with
82 additions
and
16 deletions.
There are no files selected for viewing
This file contains 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
This file contains 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