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.