diffrax
Optimization.jl
Our great sponsors
diffrax | Optimization.jl | |
---|---|---|
21 | 3 | |
1,230 | 663 | |
- | 3.3% | |
8.3 | 9.7 | |
6 days ago | about 7 hours ago | |
Python | Julia | |
Apache License 2.0 | MIT License |
Stars - the number of stars that a project has on GitHub. Growth - month over month growth in stars.
Activity is a relative number indicating how actively a project is being developed. Recent commits have higher weight than older ones.
For example, an activity of 9.0 indicates that a project is amongst the top 10% of the most actively developed projects that we are tracking.
diffrax
- Ask HN: What side projects landed you a job?
-
[P] Optimistix, nonlinear optimisation in JAX+Equinox!
Optimistix has high-level APIs for minimisation, least-squares, root-finding, and fixed-point iteration and was written to take care of these kinds of subroutines in Diffrax.
-
Show HN: Optimistix: Nonlinear Optimisation in Jax+Equinox
Diffrax (https://github.com/patrick-kidger/diffrax).
Here is the GitHub: https://github.com/patrick-kidger/optimistix
The elevator pitch is Optimistix is really fast, especially to compile. It
-
Scientific computing in JAX
Sure. So I've got some PyTorch benchmarks here. The main take-away so far has been that for a neural ODE, the backward pass takes about 50% longer in PyTorch, and the forward (inference) pass takes an incredible 100x longer.
-
[D] JAX vs PyTorch in 2023
FWIW this worked for me. :D My full-time job is now writing JAX libraries at Google. Equinox for neural networks, Diffrax for differential equation solvers, etc.
-
Returning to snake's nest after a long journey, any major advances in python for science ?
It's relatively early days yet, but JAX is in the process of developing its nascent scientific computing / scientific machine learning ecosystem. Mostly because of its strong autodifferentiation capabilities, excellent JIT compiler etc. (E.g. to show off one of my own projects, Diffrax is the library of diffeq solvers for JAX.)
-
What's the best thing/library you learned this year ?
Diffrax - solving ODEs with Jax and computing it's derivatives automatically functools - love partial and lru_cache fastprogress - simpler progress bar than tqdm
-
PyTorch 2.0
At least prior to this announcement: JAX was much faster than PyTorch for differentiable physics. (Better JIT compiler; reduced Python-level overhead.)
E.g for numerical ODE simulation, I've found that Diffrax (https://github.com/patrick-kidger/diffrax) is ~100 times faster than torchdiffeq on the forward pass. The backward pass is much closer, and for this Diffrax is about 1.5 times faster.
It remains to be seen how PyTorch 2.0 will compare, or course!
Right now my job is actually building out the scientific computing ecosystem in JAX, so feel free to ping me with any other questions.
-
Python 3.11 is much faster than 3.8
https://github.com/patrick-kidger/diffrax
Which are neural network and differential equation libraries for JAX.
[Obligatory I-am-googler-my-opinions-do-not-represent- your-employer...]
-
Ask HN: What's your favorite programmer niche?
Autodifferentiable programming!
Neural networks are the famous example of this, of course -- but this can be extended to all of scientific computing. ODE/SDE solvers, root-finding algorithms, LQP, molecular dynamics, ...
These days I'm doing all my work in JAX. (E.g. see Equinox or Diffrax: https://github.com/patrick-kidger/equinox, https://github.com/patrick-kidger/diffrax). A lot of modern work is now based around hybridising such techniques with neural networks.
I'd really encourage anyone interested to learn how JAX works under-the-hood as well. (Look up "autodidax") Lots of clever/novel ideas in its design.
Optimization.jl
-
SciPy: Interested in adopting PRIMA, but little appetite for more Fortran code
Interesting response. I develop the Julia SciML organization https://sciml.ai/ and we'd be more than happy to work with you to get wrappers for PRIMA into Optimization.jl's general interface (https://docs.sciml.ai/Optimization/stable/). Please get in touch and we can figure out how to set this all up. I personally would be curious to try this out and do some benchmarks against nlopt methods.
-
Help me to choose an optimization framework for my problem
There are also Optimization and Nonconvex , which seem like umbrella packages and I am not sure what methods to use inside these packages. Any help on these?
-
The Julia language has a number of correctness flaws
> but would you say most packages follow or enforce SemVer?
The package ecosystem pretty much requires SemVer. If you just say `PackageX = "1"` inside of a Project.toml [compat], then it will assume SemVer, i.e. any version 1.x is non-breaking an thus allowed, but not version 2. Some (but very few) packages do `PackageX = ">=1"`, so you could say Julia doesn't force SemVar (because a package can say that it explicitly believes it's compatible with all future versions), but of course that's nonsense and there will always be some bad actors around. So then:
> Would enforcing a stricter dependency graph fix some of the foot guns of using packages or would that limit composability of packages too much?
That's not the issue. As above, the dependency graphs are very strict. The issue is always at the periphery (for any package ecosystem really). In Julia, one thing that can amplify it is the fact that Requires.jl, the hacky conditional dependency system that is very not recommended for many reasons, cannot specify version requirements on conditional dependencies. I find this to be the root cause of most issues in the "flow" of the package development ecosystem. Most packages are okay, but then oh, I don't want to depend on CUDA for this feature, so a little bit of Requires.jl here, and oh let me do a small hack for OffSetArrays. And now these little hacky features on the edge are both less tested and not well versioned.
Thankfully there's a better way to do it by using multi-package repositories with subpackages. For example, https://github.com/SciML/GalacticOptim.jl is a global interface for lots of different optimization libraries, and you can see all of the different subpackages here https://github.com/SciML/GalacticOptim.jl/tree/master/lib. This lets there be a GalacticOptim and then a GalacticBBO package, each with versioning, but with tests being different while allowing easy co-development of the parts. Very few packages in the Julia ecosystem actually use this (I only know of one other package in Julia making use of this) because the tooling only recently was able to support it, but this is how a lot of packages should be going.
The upside too is that Requires.jl optional dependency handling is by far and away the main source of loading time issues in Julia (because it blocks precompilation in many ways). So it's really killing two birds with one stone: decreasing package load times by about 99% (that's not even a joke, it's the huge majority of the time for most packages which are not StaticArrays.jl) while making version dependencies stricter. And now you know what I'm doing this week and what the next blog post will be on haha.
What are some alternatives?
deepxde - A library for scientific machine learning and physics-informed learning
StatsBase.jl - Basic statistics for Julia
tiny-cuda-nn - Lightning fast C++/CUDA neural network framework
Petalisp - Elegant High Performance Computing
flax - Flax is a neural network library for JAX that is designed for flexibility.
OffsetArrays.jl - Fortran-like arrays with arbitrary, zero or negative starting indices.
juliaup - Julia installer and version multiplexer
avm - Efficient and expressive arrayed vector math library with multi-threading and CUDA support in Common Lisp.
equinox - Elegant easy-to-use neural networks + scientific computing in JAX. https://docs.kidger.site/equinox/
Distributions.jl - A Julia package for probability distributions and associated functions.
dm-haiku - JAX-based neural network library
StaticLint.jl - Static Code Analysis for Julia