Fix PARSeq decoder for TFLite compatibility #2467
Conversation
Use ops.cond() instead of Python if-statement in PARSeqDecoder.call() to ensure graph mode compatibility during TFLite conversion. - Changed 'if tokens_length > 1:' to ops.cond(tokens_length > 1, ...) - This allows TensorFlow to properly trace the graph with symbolic tensors This fixes the 'Using a symbolic tf.Tensor as a Python bool is not allowed' error during TFLite conversion when sequence length is dynamic (None).
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 focuses on enhancing the robustness and compatibility of the PARSeq decoder, specifically for TensorFlow Lite and general graph-mode execution. By refactoring a key conditional block in the Highlights
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 correctly refactors the call method in parseq_decoder.py to use ops.cond instead of a standard Python if statement. This is a crucial change for ensuring compatibility with TensorFlow Lite and graph mode execution, where control flow must be handled within the computation graph. The implementation is clean and correct, defining a separate function for the conditional logic, which improves readability. This change effectively addresses the issue and adheres to best practices for writing backend-agnostic and graph-compatible Keras code.
|
@gemini-code-assist review |
There was a problem hiding this comment.
Code Review
This pull request updates the PARSeqDecoder to be compatible with TensorFlow Lite and graph mode by replacing a Python conditional with a branchless tensor-based implementation. The approach is generally correct and improves backend compatibility. However, I've identified an edge case in the new implementation where an input with a sequence length of 0 would cause incorrect slicing and likely lead to a runtime error. My review includes a specific code suggestion to fix this issue.
pctablet505
requested review from
sachinprasadhs
and removed request for
sachinprasadhs
Simplifies content and query embedding construction for better compatibility with JAX/TF graph backends. Removes dynamic slicing and Python conditionals, using ops.take and shape-based indexing to ensure consistent tensor shapes.
3 participants
This pull request updates the
callmethod inparseq_decoder.pyto improve compatibility with TensorFlow Lite and graph mode. The main change is replacing a standard Python conditional withops.cond, ensuring the code works correctly in both eager and graph execution environments.TensorFlow Lite/graph mode compatibility:
if/elselogic for building thecontenttensor with anops.condcall, enabling proper conditional execution in graph mode. This change moves the content construction logic into a function and usesops.condto select between two execution paths based on the value oftokens_length.