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

Coherent policy needed for use of jax.jit #323

Open
bwohlberg opened this issue Aug 8, 2022 · 2 comments
Open

Coherent policy needed for use of jax.jit #323

bwohlberg opened this issue Aug 8, 2022 · 2 comments
Assignees
@bwohlberg @FernandoDavis @Michael-T-McCann
Labels
discussion required Some discussion necessary to decide how to address this issue improvement Improvement of existing code, including addressing of omissions or inconsistencies

Comments

@bwohlberg
Copy link
Collaborator

bwohlberg commented Aug 8, 2022 •
edited by Michael-T-McCann

A coherent policy for the use of jax.jit needs to be formulated and implemented. At present, its usage seems somewhat arbitrary and inconsistent; for example:

  1. Operator defaults to jit=False while LinearOperator defaults to jit=True
  2. For the most part there is no jitting of Functional classes, with the exception of L0Norm.prox.
  3. Refactoring of methods in scico.optimize classes to support jitting should be explored.
@bwohlberg bwohlberg added discussion required Some discussion necessary to decide how to address this issue improvement Improvement of existing code, including addressing of omissions or inconsistencies labels Aug 8, 2022
@bwohlberg bwohlberg added this to the Release 0.0.3 milestone Aug 8, 2022
@Michael-T-McCann Michael-T-McCann mentioned this issue Sep 6, 2022
Closed
@bwohlberg bwohlberg removed this from the Release 0.0.3 milestone Sep 12, 2022
@bwohlberg
Copy link
Collaborator Author

A suitable place should be found in the docs for the jax notes on the use of jit.

@bwohlberg
Copy link
Collaborator Author

Suggestion for a coherent policy made during a discussion of #467: all functions that can be jitted should be jitted at the point of definition using the @jax.jit decorator.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Assignees

@bwohlberg bwohlberg

@FernandoDavis FernandoDavis

@Michael-T-McCann Michael-T-McCann

Labels
discussion required Some discussion necessary to decide how to address this issue improvement Improvement of existing code, including addressing of omissions or inconsistencies
Projects
None yet
Milestone
No milestone
Development

No branches or pull requests
3 participants
@bwohlberg @FernandoDavis @Michael-T-McCann