Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Add jax.jit(donate_argnames=...) #10539

Closed
PhilipVinc opened this issue May 3, 2022 · 2 comments
Closed

Add jax.jit(donate_argnames=...) #10539

PhilipVinc opened this issue May 3, 2022 · 2 comments
Assignees
@yashk2810
Labels
enhancement New feature or request

Comments

@PhilipVinc
Copy link
Contributor

static_argnames is much less error prone than static_argnums, and easily works with keyword arguments.
It would be useful to have the same syntax for donate_***.
@PhilipVinc PhilipVinc added the enhancement New feature or request label May 3, 2022
@JeppeKlitgaard
Copy link
Contributor

I might be interested in implementing this following #10603 and #10476

@JeppeKlitgaard
Copy link
Contributor

Having had a look at this, I think an ideal solution which supports both positional-only and keyword-only arguments (as well as anything in between) would require changes to xla.xla_call, which would be beyond what I could do.

A simpler solution that would work in cases except keyword-only could be done quite easily with inspect.

@JeppeKlitgaard JeppeKlitgaard mentioned this issue May 6, 2022
Open
15 tasks
@yashk2810 yashk2810 self-assigned this Feb 7, 2023
@copybara-service copybara-service bot mentioned this issue Jul 7, 2023
Closed
copybara-service bot pushed a commit that referenced this issue Jul 7, 2023
@yashk2810
Add donate_argnames to jax.jit. This acts similarly to static_argnames.
dce2203 
Fixes #10539

PiperOrigin-RevId: 546125263
copybara-service bot pushed a commit that referenced this issue Jul 7, 2023
@yashk2810
Add donate_argnames to jax.jit. This acts similarly to static_argnames.
500d243 
Fixes #10539

PiperOrigin-RevId: 546125263
copybara-service bot pushed a commit that referenced this issue Jul 11, 2023
@yashk2810
Add donate_argnames to jax.jit. This works similarly to static_argnames.
a946973 
Note that if donate_argnames is not None and donate_argnums is None, then JAX will infer donate_argnums from the names which will then we used to find the donation_vector. This is fine because currently, the same thing happens from static_argnums and static_argnames.

I'll fix the TODOs, etc in follow up CLs.

Fixes #10539

PiperOrigin-RevId: 546125263
copybara-service bot pushed a commit that referenced this issue Jul 11, 2023
@yashk2810
Add donate_argnames to jax.jit. This works similarly to static_argnames.
f38e96d 
Note that if donate_argnames is not None and donate_argnums is None, then JAX will infer donate_argnums from the names which will then we used to find the donation_vector. This is fine because currently, the same thing happens from static_argnums and static_argnames.

I'll fix the TODOs, etc in follow up CLs.

Fixes #10539

PiperOrigin-RevId: 546125263
copybara-service bot pushed a commit that referenced this issue Jul 11, 2023
@yashk2810
Add donate_argnames to jax.jit. This works similarly to static_argnames.
625d901 
Note that if donate_argnames is not None and donate_argnums is None, then JAX will infer donate_argnums from the names which will then we used to find the donation_vector. This is fine because currently, the same thing happens from static_argnums and static_argnames.

I'll fix the TODOs, etc in follow up CLs.

Fixes #10539

PiperOrigin-RevId: 546125263
copybara-service bot pushed a commit that referenced this issue Jul 12, 2023
@yashk2810
Add donate_argnames to jax.jit. This works similarly to static_argnames.
cbc48a0 
Note that if donate_argnames is not None and donate_argnums is None, then JAX will infer donate_argnums from the names which will then we used to find the donation_vector. This is fine because currently, the same thing happens from static_argnums and static_argnames.

I'll fix the TODOs, etc in follow up CLs.

Fixes #10539

PiperOrigin-RevId: 546125263
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Assignees

@yashk2810 yashk2810

Labels
enhancement New feature or request
Projects
None yet
Milestone
No milestone
Development

Successfully merging a pull request may close this issue.

Add donate_argnames to jax.jit. This works similarly to static_argnames. google/jax
3 participants
@JeppeKlitgaard @PhilipVinc @yashk2810