Skip to content

optimize jit_partial - #5476

Merged
copybara-service[bot] merged 1 commit into
mainfrom
test_923984648
Jun 17, 2026
Merged

copybara-service[bot] merged 1 commit into
mainfrom
test_923984648

Conversation

@copybara-service

@copybara-service copybara-service Bot commented May 31, 2026 •

Copy link
Copy Markdown

optimize jit_partial

Optimizes jit_partial by various means:

  • Creates specialized call path for jit_partial (SimpleJitPartialWrapped and friends).
  • Flattens the partials structure ahead of time to minimize traversals during call.
  • Forbids Variables in non-partial arguments to avoid alias checks.
  • Adds specialized _apply_raw_updates helper which optimizes trace context checks.
  • Updates only contain inner values and metadata updates not allowed (could be lifted in the future).

@copybara-service
copybara-service Bot force-pushed the test_923984648 branch 5 times, most recently from 9995c47 to 796febc Compare June 17, 2026 06:44
Optimizes jit_partial by various means:
* Creates specialized call path for jit_partial (SimpleJitPartialWrapped and friends).
* Flattens the partials structure ahead of time to minimize traversals during __call__.
* Forbids Variables in non-partial arguments to avoid alias checks.
* Adds specialized _apply_raw_updates helper which optimizes trace context checks.
* Updates only contain inner values and metadata updates not allowed (could be lifted in the future).

PiperOrigin-RevId: 933945334
@copybara-service
copybara-service Bot merged commit 0a0aca7 into main Jun 17, 2026
@copybara-service
copybara-service Bot deleted the test_923984648 branch June 17, 2026 22:11
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

0 participants