⚡ Optimize StateObserver to avoid round-trip state conversion#9
⚡ Optimize StateObserver to avoid round-trip state conversion#9
Conversation
Co-authored-by: google-labs-jules[bot] <161369871+google-labs-jules[bot]@users.noreply.github.com>
|
👋 Jules, reporting for duty! I'm here to lend a hand with this pull request. When you start a review, I'll add a 👀 emoji to each comment to let you know I've read it. I'll focus on feedback directed at me and will do my best to stay out of conversations between you and other bots or reviewers to keep the noise down. I'll push a commit with your requested changes shortly after. Please note there might be a delay between these steps, but rest assured I'm on the job! For more direct control, you can switch me to Reactive Mode. When this mode is on, I will only act on comments where you specifically mention me with New to Jules? Learn more at jules.google/docs. For security, I will only act on instructions from the user who triggered this task. |
Summary of ChangesHello @igor-holt, 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 significantly enhances the performance of the Highlights
Changelog
Activity
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 a valuable performance optimization to StateObserver.__call__ by avoiding an expensive round-trip state conversion. The logic is sound and the inclusion of benchmarks is great. I have one suggestion to refactor a part of the new logic to improve code clarity.
| sampled_state = [None] * len(self.blocks_to_sample) | ||
| blocks_missing = [] | ||
| indices_missing = [] | ||
|
|
||
| for i, block in enumerate(self.blocks_to_sample): | ||
| if block in block_map: | ||
| source, idx = block_map[block] | ||
| sampled_state[i] = source[idx] | ||
| else: | ||
| blocks_missing.append(block) | ||
| indices_missing.append(i) | ||
|
|
||
| if blocks_missing: | ||
| global_state = block_state_to_global(state_free + state_clamped, program.gibbs_spec) | ||
| missing_sampled = from_global_state(global_state, program.gibbs_spec, blocks_missing) | ||
| for idx, s in zip(indices_missing, missing_sampled): | ||
| sampled_state[idx] = s |
There was a problem hiding this comment.
To improve clarity and maintainability, you could refactor the logic for handling missing blocks. Instead of managing two separate lists, blocks_missing and indices_missing, you could use a single list of tuples (index, block). This is a common Python pattern that can make the code cleaner and less error-prone, as it ensures the indices and blocks always stay in sync. You can then use zip(*missing) to unzip this list back into two separate sequences when needed.
| sampled_state = [None] * len(self.blocks_to_sample) | |
| blocks_missing = [] | |
| indices_missing = [] | |
| for i, block in enumerate(self.blocks_to_sample): | |
| if block in block_map: | |
| source, idx = block_map[block] | |
| sampled_state[i] = source[idx] | |
| else: | |
| blocks_missing.append(block) | |
| indices_missing.append(i) | |
| if blocks_missing: | |
| global_state = block_state_to_global(state_free + state_clamped, program.gibbs_spec) | |
| missing_sampled = from_global_state(global_state, program.gibbs_spec, blocks_missing) | |
| for idx, s in zip(indices_missing, missing_sampled): | |
| sampled_state[idx] = s | |
| sampled_state = [None] * len(self.blocks_to_sample) | |
| missing = [] | |
| for i, block in enumerate(self.blocks_to_sample): | |
| if block in block_map: | |
| source, idx = block_map[block] | |
| sampled_state[i] = source[idx] | |
| else: | |
| missing.append((i, block)) | |
| if missing: | |
| indices_missing, blocks_missing = zip(*missing) | |
| global_state = block_state_to_global(state_free + state_clamped, program.gibbs_spec) | |
| missing_sampled = from_global_state( | |
| global_state, program.gibbs_spec, list(blocks_missing) | |
| ) | |
| for idx, s in zip(indices_missing, missing_sampled): | |
| sampled_state[idx] = s |
1 participant
💡 What:
Optimized
StateObserver.__call__inthrml/observers.pyto check if the requested blocks are directly available instate_freeorstate_clampedbefore falling back toblock_state_to_global.🎯 Why:
The previous implementation always converted the entire block state to a global vector and then extracted the specific blocks using
from_global_state. This involved unnecessary data copying and overhead, especially when most blocks are already available in the input state lists. The round-trip conversion was a performance bottleneck.📊 Measured Improvement:
A benchmark (
benchmark_observer.py) was created to measure the execution time ofStateObserver.__call__.The optimization logic handles mixed cases correctly (where some blocks are available directly and others need reconstruction) and falls back to the original method when necessary.
Verified correctness with a new test case
tests/test_state_observer_optimization.pyand existing teststests/test_observers.pyandtests/test_block_sampling.py.PR created automatically by Jules for task 677721227058151632 started by @igor-holt