Comments (8)
At the moment I essentially use jaxopt with ed1febe reverted and make sure that I cache my solvers:
@lru_cache
def get_solver(solver, *args, **kwargs):
return solver(*args, **kwargs)
My suggestion would be to
- Implement caching of
_while_loop_lax
cleanly without relying onjax.jit
-- I haven't dug deep enough into jax yet to know the best way to do this, and - Make it easy for the user to reuse solvers or at least document that reusing solvers will reduce / avoid recompilation to benefit from this change. Another topic would be advising against nested functions, which I've seen in a lot of non-official examples.
from jaxopt.
CC @froystig
from jaxopt.
This is mostly relevant when reusing the same solver (I've been using BFGS), as reinstantiation creates new references of _cond_fun
and _body_fun
, which are static arguments in
Line 80 in 58bac0a
from jaxopt.
Do you have a minimal example that reproduces the slowdown, and would you mind posting it here if so?
from jaxopt.
import time, jaxopt, jax, jax.numpy as jnp
def rosenbrock(x):
return jnp.sum(100. * jnp.diff(x) ** 2 + (1. - x[:-1]) ** 2)
solver = jaxopt.BFGS(rosenbrock)
x0 = jnp.zeros(2)
_time = time.time()
sol = solver.run(x0)
_time = time.time() - _time
print(f'Total {_time} s')
jax.config.update('jax_log_compiles', True)
_time = time.time()
sol2 = solver.run(x0)
_time = time.time() - _time
print(f'Total2 {_time} s')
With jax.jit
:
Total 1.3411040306091309 s
[1. 1.]
Total2 0.007392406463623047 s
Original library (jax 0.4.23, jaxopt 0.8.2):
Total 1.3537604808807373 s
Finished tracing + transforming while for pjit in 0.0003256797790527344 sec
Compiling while for with global shapes and types [ShapedArray(int32[], weak_type=True), ShapedArray(float32[2]), ShapedArray(int32[], weak_type=True), ShapedArray(float32[]), ShapedArray(float32[2]), ShapedArray(float32[]), ShapedArray(float32[]), ShapedArray(float32[2,2]), ShapedArray(int32[]), ShapedArray(int32[]), ShapedArray(int32[])]. Argument mapping: (GSPMDSharding({replicated}), GSPMDSharding({replicated}), GSPMDSharding({replicated}), GSPMDSharding({replicated}), GSPMDSharding({replicated}), GSPMDSharding({replicated}), GSPMDSharding({replicated}), GSPMDSharding({replicated}), GSPMDSharding({replicated}), GSPMDSharding({replicated}), GSPMDSharding({replicated})).
Finished jaxpr to MLIR module conversion jit(while) in 0.1349470615386963 sec
Finished XLA compilation of jit(while) in 0.3022458553314209 sec
[1. 1.]
Total2 0.45375514030456543 s
from jaxopt.
More generally, what do you think about jaxopt caching all solvers, so recompilation would be reduced automatically when not using nested functions?
from jaxopt.
what do you propose more concretely?
from jaxopt.
Would it be possible to detect that you're inside a jitted context rather than accepting the Boolean jit
parameter? That way, only the user would ever call jit
, and would totally control caching and compilation.
from jaxopt.
Related Issues (20)
- Consider switching to pyproject.toml
- OSQP crashing on unexpected params HOT 3
- `verbose=False` is not working as expected for `NonlinearCG` HOT 1
- Stochastic L-BFGS algorithm implementation
- Stopping condition 'madsen-nielsen' incorrect
- unit test failures on aarch64 linux with scipy 1.12
- `LevenbergMarquardt` implementation does not accept PyTree parameters
- diag(JTJ) can be more efficient
- JAXOPT Projected Gradient
- "invalid escape sequence" warning in `BoxOSQP` docstring
- Error when taking gradient wrt parameters in BoxOSQP
- Disable warnings in vmap or print to stderr HOT 2
- Constraint violation causes L-BGFS-B to fail HOT 1
- LinearSolveTest.test_solve_sparse fails with jax 0.4.26 HOT 1
- Tests crash/failures on `aarch64-linux` with latest jax(lib) (0.4.28)
- Projected gradient for multidimensional array
- Annoying warning in 3.12
- pytrees bounds for `jaxopt.ScipyBoundedMinimize` HOT 1
- Jaxopt implementation of scipy's linprog
- GaussNewton and LevenbergMarquardt on multidimensional arrays
Recommend Projects
-
React
A declarative, efficient, and flexible JavaScript library for building user interfaces.
-
Vue.js
🖖 Vue.js is a progressive, incrementally-adoptable JavaScript framework for building UI on the web.
-
Typescript
TypeScript is a superset of JavaScript that compiles to clean JavaScript output.
-
TensorFlow
An Open Source Machine Learning Framework for Everyone
-
Django
The Web framework for perfectionists with deadlines.
-
Laravel
A PHP framework for web artisans
-
D3
Bring data to life with SVG, Canvas and HTML. 📊📈🎉
-
Recommend Topics
-
javascript
JavaScript (JS) is a lightweight interpreted programming language with first-class functions.
-
web
Some thing interesting about web. New door for the world.
-
server
A server is a program made to process requests and deliver data to clients.
-
Machine learning
Machine learning is a way of modeling and interpreting data that allows a piece of software to respond intelligently.
-
Visualization
Some thing interesting about visualization, use data art
-
Game
Some thing interesting about game, make everyone happy.
Recommend Org
-
Facebook
We are working to build community through open source technology. NB: members must have two-factor auth.
-
Microsoft
Open source projects and samples from Microsoft.
-
Google
Google ❤️ Open Source for everyone.
-
Alibaba
Alibaba Open Source for everyone
-
D3
Data-Driven Documents codes.
-
Tencent
China tencent open source team.
from jaxopt.