Skip to content

Support unused arguments in xb.call_jax#8830

Merged
tengyifei merged 1 commit intomasterfrom
yifeit/call-jax-args
Mar 13, 2025
Merged

Support unused arguments in xb.call_jax#8830
tengyifei merged 1 commit intomasterfrom
yifeit/call-jax-args

Conversation

@tengyifei
Copy link
Copy Markdown
Collaborator

@tengyifei tengyifei commented Mar 13, 2025

Fixes #8794.

Turns out the answer is as simple as keep_unused=True.

I also took a test case in
#8800 and made it pass.

Fixes #8794.

I also took a test case in
#8800 and made it pass.
@tengyifei tengyifei marked this pull request as ready for review March 13, 2025 01:41
@tengyifei tengyifei requested review from qihqi and zpcore March 13, 2025 01:45
Copy link
Copy Markdown
Member

@zpcore zpcore left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is nice!

@tengyifei tengyifei merged commit b3f3875 into master Mar 13, 2025
zpcore pushed a commit that referenced this pull request Mar 26, 2025
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.

xb.call_jax always assumes the traced function makes use of all arguments

2 participants