[catgary]

The ODEs stuff in Julia is nice, but I think diffusers/JAX is a reasonable backbone to copy over whatever you need from there. I do think Julia is doing will in stats and has gotten some mindshare from R in that regard.

But I think a reasonably compentent Python/JAX programmer can roll out whatever they need relatively easily (especially if you want to use the GPU). I do miss Tullio, though.

[ssivark]

But how does one get the composability of multiple dispatch? A problem with the Python world IME is the need to reinvent everything in Jax ecosystem, in the Pytorch ecosystem, etc -- because of fundamental language limitations. This is tedious, and feels like a lot of wasted effort.

Another example: It's frustrating that Flax had to implement it's own "lifted" transformations instead of being able to just use jax transformations -- which makes it impossible to just slot a Flax model into a jax library that integrates ODEs. Equinox might be better on this front, but that means that all the models now need to be re-implemented in Equinox. The fragmentation and churn in the Python ecosystem is outrageous -- the only reason it doesn't collapse under its own weight is how much funding and manpower ML stakeholders are able to pour into the ecosystem.

Given how much the ecosystem depends on that sponsored effort, the popular frameworks will likely prioritize ML applications, and corollary use cases will be second class citizens in case of design tradeoffs. Eg: framework overheads matter less when one is trying to use large NN models -vs- when one is trying to use small models, or other parametric approaches.

[catgary]

So, I’ll admit that I’m not a fan of multiple dispatch a la Julia. I much prefer typeclasses and explicit union types. Also, I found Julia is really the worst of both worlds of garbage collected languages and manual memory management because it ostensibly has a garbage collector but you find yourself basic preallocating memory or faffing around with StaticArrays trying to figure out why memory isn’t being allocated (often it comes down to some type instability nonsense because the type system built around multiple dispatch can’t correctly type the program). At this point I’d rather just use C++ or Rust than Julia, I’m getting annoyed just thinking about the nonsense I used to deal with.

Also, IIRC, it’s not terribly difficult to use flax with equinox. It’s just a matter of storing the weight dict and model function in an equinox module. Filter_jit will correctly recognize the weights as a dynamic variable and the flax model as a static variable.

[lagrange77]

> But I think a reasonably compentent Python/JAX programmer can roll out whatever they need relatively easily

You mean in terms of the ODE stuff, Julia provides?

[catgary]

Diffusers is pretty well done (I think the author was basically rewriting some Julia libraries and adapting them to JAX). I can’t imagine it being too hard to adapt most SciML ODE solvers.

For simulations, JAX will choke on very “branchy” computations. But, honestly I’ve had very little success differentiating through those computations in the first place and they don’t run well on the GPU. Thus, I’m generally inclined to use wrappers around C++ (or ideally Rust) for those purposes (my use-case is usually some rigid-body dynamics style simulation).