extending-jax VS diffrax

Compare extending-jax vs diffrax and see what are their differences.

extending-jax

Extending JAX with custom C++ and CUDA code (by dfm)

diffrax

Numerical differential equation solvers in JAX. Autodifferentiable and GPU-capable. https://docs.kidger.site/diffrax/ (by patrick-kidger)
Scout Monitoring - Free Django app performance insights with Scout Monitoring
Get Scout setup in minutes, and let us sweat the small stuff. A couple lines in settings.py is all you need to start monitoring your apps. Sign up for our free tier today.
www.scoutapm.com
featured
InfluxDB - Power Real-Time Data Analytics at Scale
Get real-time insights from all types of time series data with InfluxDB. Ingest, query, and analyze billions of data points in real-time with unbounded cardinality.
www.influxdata.com
featured
extending-jax diffrax
2 21
361 1,324
- -
3.5 8.6
8 months ago 5 days ago
Python Python
MIT License Apache License 2.0
The number of mentions indicates the total number of mentions that we've tracked plus the number of user suggested alternatives.
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.

extending-jax

Posts with mentions or reviews of extending-jax. We have used some of these posts to build our list of alternatives and similar projects. The last one was on 2022-02-15.

diffrax

Posts with mentions or reviews of diffrax. We have used some of these posts to build our list of alternatives and similar projects. The last one was on 2023-12-03.
  • Ask HN: What side projects landed you a job?
    62 projects | news.ycombinator.com | 3 Dec 2023
  • [P] Optimistix, nonlinear optimisation in JAX+Equinox!
    3 projects | /r/MachineLearning | 14 Oct 2023
    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
    2 projects | news.ycombinator.com | 10 Oct 2023
    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
    4 projects | /r/ScientificComputing | 4 Apr 2023
    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
    5 projects | /r/MachineLearning | 9 Mar 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 ?
    7 projects | /r/Python | 24 Jan 2023
    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 ?
    12 projects | /r/Python | 16 Dec 2022
    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
    4 projects | news.ycombinator.com | 2 Dec 2022
    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
    11 projects | news.ycombinator.com | 26 Oct 2022
    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?
    8 projects | news.ycombinator.com | 15 Oct 2022
    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.

What are some alternatives?

When comparing extending-jax and diffrax you can also consider the following projects:

einops - Flexible and powerful tensor operations for readable and reliable code (for pytorch, jax, TF and others)

deepxde - A library for scientific machine learning and physics-informed learning

mpi4jax - Zero-copy MPI communication of JAX arrays, for turbo-charged HPC applications in Python :zap:

tiny-cuda-nn - Lightning fast C++/CUDA neural network framework

thinc - 🔮 A refreshing functional take on deep learning, compatible with your favorite libraries

juliaup - Julia installer and version multiplexer

equinox - Elegant easy-to-use neural networks + scientific computing in JAX. https://docs.kidger.site/equinox/

flax - Flax is a neural network library for JAX that is designed for flexibility.

trax - Trax — Deep Learning with Clear Code and Speed

elegy - A High Level API for Deep Learning in JAX

vectorflow

Scout Monitoring - Free Django app performance insights with Scout Monitoring
Get Scout setup in minutes, and let us sweat the small stuff. A couple lines in settings.py is all you need to start monitoring your apps. Sign up for our free tier today.
www.scoutapm.com
featured
InfluxDB - Power Real-Time Data Analytics at Scale
Get real-time insights from all types of time series data with InfluxDB. Ingest, query, and analyze billions of data points in real-time with unbounded cardinality.
www.influxdata.com
featured