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

Support symbolic zero operations in custom_jvp #5913

Open
srvasude opened this issue Mar 3, 2021 · 0 comments
Open

Support symbolic zero operations in custom_jvp #5913

srvasude opened this issue Mar 3, 2021 · 0 comments
Assignees
@froystig
Labels
enhancement New feature or request P2 (eventual) This ought to be addressed, but has no schedule at the moment. (Assignee optional)

Comments

@srvasude
Copy link
Member

srvasude commented Mar 3, 2021

It would be nice if jax.custom_jvp supported symbolic zero operations much like the deprecated defvjp functions:

https://cs.opensource.google/jax/jax/+/master:jax/api.py;drc=654a5b332c0d2f7a77b275bf31e186cfe528cdb2;l=2565

Specifically, given a bivariate function f(x, y), I would like to define the jvp of f, such that I define the jvp with respect to x, and raise an error when differentiating with respect to y. This can happen if I don't have a closed form for the jvp with respect to y, and I want users to avoid this as an issue.

In the autodiff tracer API this is allowed (for instance see betainc: https://cs.opensource.google/jax/jax/+/master:jax/_src/lax/lax.py;drc=0dd1b5516d4028cd182d01865279dcebf2f27c40;l=2372). Taking grad or jvp with respect to the x parameter of betainc gives the right answer, whereas writing a custom_jvp for betainc will throw an exception.
@srvasude srvasude added the enhancement New feature or request label Mar 3, 2021
@froystig froystig self-assigned this Mar 11, 2021
@froystig froystig added the open Issues intentionally left open, with no schedule for next steps. label Mar 11, 2021
@ahoenselaar ahoenselaar mentioned this issue Apr 12, 2021
Closed
@froystig froystig mentioned this issue Jul 1, 2022
Closed
@froystig froystig added P2 (eventual) This ought to be addressed, but has no schedule at the moment. (Assignee optional) and removed open Issues intentionally left open, with no schedule for next steps. labels Jul 1, 2022
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Assignees

@froystig froystig

Labels
enhancement New feature or request P2 (eventual) This ought to be addressed, but has no schedule at the moment. (Assignee optional)
Projects
None yet
Milestone
No milestone
Development

No branches or pull requests
2 participants
@froystig @srvasude