patrickkidger

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, of 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.

+1 for JAX. Basically designed to be the successor to TensorFlow, and much nicer to work with. Strangely I've not seen it discussed around HN much but it's what I do 100% of my work in these days.

Whilst I'm here: shameless self-promotion for Equinox and Diffrax:

https://github.com/patrick-kidger/equinox https://github.com/patrick-kidger/diffrax

Which are neural network and differential equation libraries for JAX.

[Obligatory I-am-a-googler-my-opinions-do-not-represent-your-employer...]

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.

There's lots of good reasons. Most notably the autodiff in Julia is unreliable -- there's been quite the litany of packages trying/failing to do this robustly and efficiently. (This issue pretty much kills its possibilities straight from the get-go.)

The fact that the ecosystem is largely created by academics is also a serious issue. Code quality is very low / things silently break / etc. This is partly an issue with Julia itself; the language offers few tools to verify the correctness of your code. (E.g. interfaces etc.)

Julia vs XYZ is an old topic. And yet this keeps coming up! There's quite a lot of blog posts on the topic. (E.g. here's mine: https://kidger.site/thoughts/jax-vs-julia/)

From personal experience: possibly Twitter. Having a presence there has given me several serious offers to interview, for jobs worth having (FAANG and similar).

The folks in my field (ML/applied math/scientific computing) are all on there, and moreover those reaching out to me are always technical folks rather than recruiters.

(Twitter gets a bad rap, but the academic/dev parts of it are the single best way I've found to track the Zeitgeist of fields and topics I'm interested in.)

Personally I'm a big fan of Julia's syntax here. Indentation-insensitive with the use of the word "end" as the terminator for all types of block. In practice you can lay it out like Python, and the "end"s help you keep track when e.g. you need to place a single expression after some deeply-nested blocks.

I disagree. I've had a serious attempt at array typing using variadic generics and I'm not impressed. Python's type system has numerous issues... and now they just apply to any "ArrayWithNDimensions" type as well as any "ArrayWith2Dimensions" type.

Variadic protocols don't exist; many operations like stacking are inexpressible; the synatx is awful and verbose; etc. etc.

I've written more about this here as part of my TorchTyping project: [0]

[0] https://github.com/patrick-kidger/torchtyping/issues/37#issu...

100% this. In a discussion on "cultural correctness issues prevents me from using Julia", it's very telling that the response is "more speed!"

There's been a decent number of posts based around "Julia has these problems". And I don't think that's because the world at large has a vendetta; I think it's because the world at large desperately wants to use Julia, but struggle with hard blocks that are currently preventing adoption.

FWIW I do think there's a growing acceptance in the Julia community that these concerns are real, which is good. (See the parallel discussion on the Julia Discourse.)

Excellent! I'm very happy to take contributions generalising the tool.

FWIW my take is not that Yuri is expressing "there are too many bugs" so much as he's expressing a problem in the culture surrounding Julia itself:

> But systemic problems like this can rarely be solved from the bottom up, and my sense is that the project leadership does not agree that there is a serious correctness problem.

Concisely:

1. The ecosystem is poorly put together. (It's been produced by academics rather than professional software developers.)

2. The language provides few tools to guarantee correctness. (No static typing; no interfaces.)

Personally, what I'd love to see is one of the big tech companies come on board and just write their own ecosystem. The Julia language is amazing. The ecosystem needs to be rewritten.

You may already know of it, but if you want differential-equations-in-JAX then allow me to quickly advertise Diffrax: https://github.com/patrick-kidger/diffrax (of which I am the author, disclaimer).

Nope, this was on v0.2 of SimpleChains.jl; I did check.

Anyway, something to investigate if I ever get the time.

You're discussing switching algorithms like LSODA?

This is a really good question that I don't have a neat answer to. You can actually also hit similar issues when naively vectorising some stiff algorithms that detect when to recalculate Jacobians: at each timestep, some batch elements will want the recalculation whilst some won't.

In both cases, something like what you suggest might be a possible solution. And thankfully support for such custom batching will Soon (tm) be coming to JAX.

In the mean time, Diffrax actually side-steps the problem by not implementing those kinds of solvers. New solvers usually get added on a by-request basis, so this just hasn't been an issue yet.

So:

1. Do you know what DifferentialEquations.jl does in this scenario when running on the GPU? (On the CPU needing to sychronise batch elements usually isn't such a concern.)

2. If the price of using JAX is giving up a bit of performance in this one use case, then I think I'm okay with that. The clear contenders in this space are JAX and Julia, and right now it's a concious choice to be using JAX as a better fit for the problems I'm tackling. This comparison is something I've written about before: https://discourse.julialang.org/t/state-of-machine-learning-....

Not at all -- I was very happy to see some Equinox code out in the wild!

I think the main slowdown was doing the model updates out of the JIT region. JAX-without-JIT is ridiculously slow.

I've just run this myself.

Tidying up the JAX script, I find that it runs in ~15 seconds on my laptop CPU. Likewise, fixing the crash error in the blog post (UndefVarError: alloc_threaded_grad not defined) then I find that the Julia implementation also runs in ~15 seconds on my laptop CPU.

So, 1x faster than JAX?

Clearly we're getting very different results here, so (a) I'd like to get to the bottom of this, but (b) perhaps we should be more cautious about making performance claims.

(EDIT: updated script in this GitHub gist: https://gist.github.com/patrick-kidger/68bf7b99ba02c246b20ea...)

Taking a union of differential equations is a pretty bad idea for exactly the reasons you describe. But it's absolutely possible to parallelise multiple diffeq solves without needing to use the same time steps for each solve. That is, the time steps can be vectorised, rather than broadcast, over the batch axis.

So the only thing you actually need are the same number of steps for each solve, which can be easily accomplished by just padding out the solves that finish with slightly fewer steps. In practice this ends up introducing negligible overhead whilst solving the above issue very neatly. For example this is precisely what Diffrax (https://github.com/patrick-kidger/diffrax) does under `jax.vmap`.

I've not dug into what Julia does here; is this not already done when broadcasting `DifferentialEquations.solve`?

I'm happy to hear about Equinox being used! (I'm the author.)

I'm curious what your workloads are that you're seeing speedups of as much as 1e4? Greatest I've heard of before was ~1e2 on some differential equation solving.