Jax is awesome. But supporting most of numpy isn't enough. Because numpy isn't composable. You want to add banded-block-banded matrices to numpy? Then you need to fork numpy (or in this case fork jax); this is a package in julia and it works with everything. You want add names to your array dimentions like PyTorch recently did, then like PyTorch you need to fork numpy; again this is package in julia. You want to do both? You have to merge those two forks into each other. In julia this isn't even a package this is just using the aformentioned two packages.

You want to work with Units or track Measurment error (or both?). Basically same story. Except better in some ways worse in others. Better because you don't have to fork numpy, it is extensible enough to allow that. Packages exist that use that etendability for exactly that. Worse because those are scalar types, why are you even having to write code to deal with array support at all. Agian 2 julia packages and they don't even mention arrays internally.

The problem's not Jax. The problem is numpy. Or rather the problem is this level of composability is really hard most of the time in most languages (including the python + C combo. Especially so even).

Its true that this is not always trivial 100$% of the time with julia's multiple dispatch. but it is truer there than anywhere else i have seen.