Fix jax abstractification error #2492
Conversation
* Update gemma3_causal_lm_preprocessor.py Added checks for invalid inputs * Update gemma3_causal_lm_preprocessor.py * Update gemma3_causal_lm_preprocessor_test.py Added tests to check invalid inputs * Update reversible_embedding.py Fix for model not loading when using numpy behaviour with tensorflow * upadated Gemma3InterleaveEmbeddings * Update gemma3_interleave_embeddings.py * Revert "Update reversible_embedding.py" This reverts commit 3fdc7fd. * Update gemma3_interleave_embeddings.py * Fix ReversibleEmbedding quantization for Keras 3.13 compatibility - Update _int8_build method to accept config parameter - Ensure backward compatibility with Keras 3.12 - Resolves TypeError in quantization tests
…am#2482) * Models should reference ReversibleEmbedding from Keras core * backwards compat * remove keras 3.9 from test matrix * python version changes * disable cache when fetching nightly
Summary of ChangesHello @pctablet505, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request primarily focuses on modernizing the project's dependencies and simplifying its codebase. It updates the minimum required Python version to 3.11 and the default Keras version to 3.13, ensuring the project leverages recent language and framework improvements. A significant simplification involves delegating the Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Ignored Files
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces several important updates and fixes. The upgrade of the minimum Python version to 3.11 and Keras to 3.13 across the project ensures the use of more recent features and bug fixes. The most significant improvement is the refactoring of ReversibleEmbedding to use the upstream implementation from Keras. This greatly simplifies the codebase, reduces maintenance overhead, and aligns perfectly with the repository's style guide on reusability. The core fix for the JAX abstractification error in CausalLM by explicitly using variable.value is also a crucial change for backend compatibility.
I have a couple of minor suggestions for code cleanup in a build script and a Python file, but overall, this is a high-quality contribution that improves the project's health and stability.
.kokoro/github/ubuntu/gpu/build.sh
Outdated
| cd "${KOKORO_ROOT}/" | ||
|
|
There was a problem hiding this comment.
This cd command is redundant as the directory is already changed to ${KOKORO_ROOT}/ on line 17. You can remove this line and the following blank line for clarity.
keras_hub/src/models/causal_lm.py
Outdated
| sampler_variables = [] | ||
| for v in self.sampler.variables: | ||
| new_sampler_variables = [] | ||
| for i, v in enumerate(self.sampler.variables): |
There was a problem hiding this comment.
The loop variable i from enumerate is not used within the loop. It's cleaner to iterate directly over self.sampler.variables.
| for i, v in enumerate(self.sampler.variables): | |
| for v in self.sampler.variables: |
Eliminated a duplicate 'cd' command in the GPU build script. Refactored the sampler variable update loop in causal_lm.py to remove the unused enumerate, improving code clarity.
This reverts commit 4644ffe.
3 participants
This pull request focuses on updating the minimum Python version requirement to 3.11 across the project, updating the default Keras version in CI workflows, and simplifying the implementation of the
ReversibleEmbeddinglayer by delegating to the upstream Keras implementation. It also removes the local test suite forReversibleEmbeddingand updates documentation and environment setup instructions accordingly.Python and Keras version updates:
README.md,CONTRIBUTING.md) to require Python 3.11 or later instead of 3.10, ensuring consistency across development, testing, and deployment environments. [1] [2] [3] [4] [5] [6] [7]Environment and build script changes:
.kokoro/github/ubuntu/gpu/build.sh, setDEBIAN_FRONTEND=noninteractivefor non-interactive package installs and temporarily relaxed thepyproject.tomlPython requirement to allow Python 3.10 for GPU tests. [1] [2]Codebase simplification and maintenance:
ReversibleEmbeddinginkeras_hub/src/layers/modeling/reversible_embedding.pywith a direct subclass of the upstream Keras version, delegating all logic.token_and_position_embedding.pyto use the upstreamReversibleEmbeddingdirectly.ReversibleEmbeddingfromreversible_embedding_test.py.## Description of the changeReference
Colab Notebook
Checklist