JAX has a reputation for being hard, and the reputation is half deserved. The API is small and the individual functions are simple. What is hard is that JAX asks you to think differently about your code than numpy or PyTorch do, and if you try to write JAX the way you write PyTorch you will spend a week confused about why nothing works. The official documentation is unusually good at addressing exactly that problem.
The JAX 101 tutorial series is the entry point and it is structured as a sequence rather than a reference. It covers just in time compilation, automatic vectorisation with vmap, automatic differentiation with grad, pytrees, pseudorandom number generation, distributed arrays and sharding, control flow under jit, the tracing model, and how to handle state in a system that dislikes it. Reading it in order matters. Each section assumes the one before, and the payoff is that by the end you have a mental model rather than a list of functions.
The section on tracing is the one I would point anyone toward first, and it is the reason I rate these docs as highly as I do. Most libraries hide their execution model and let you discover it through bugs. JAX explains that when you jit a function, it is traced with abstract values, that your Python print statements will fire once during tracing and never again, that your Python conditionals will be baked in based on the traced values, and that this is by design. Getting this explained clearly before you hit it saves days.
The randomness material is a similar case. JAX's explicit PRNG keys are the single most common thing new users complain about, because it feels like unnecessary ceremony compared to setting a global seed. The docs explain the reasoning properly. Reproducibility under parallelism and compilation genuinely requires it, and once you understand why, the ceremony stops feeling arbitrary.
Explaining a design decision rather than just documenting the API is what separates good technical writing from adequate technical writing. The sharding and distributed computation coverage is stronger than most frameworks manage. JAX's approach of describing how arrays are laid out across devices and letting the compiler handle the rest is elegant, and the documentation walks through it with real examples rather than gesturing at it. If you are working on anything multi device this is a genuine advantage over the alternatives.
Now the limitations, and the first is definitional. This is documentation. It is being reviewed here as a learning resource because that is how many people use it, but there are no exercises, no problem sets, no projects and no way to check whether you understood anything. You read it, you feel like you understood it, and then you discover three weeks later that you did not.
Anyone who learns by doing needs to pair this with a project of their own invention, because the docs will not supply one. The second gap is the ecosystem question. JAX by itself is a numerical computing library with autodiff and compilation. Actually training a neural network means picking a layer library on top, and the community has produced several with different philosophies and different levels of maintenance.
The documentation is understandably reluctant to pick a winner, which leaves a genuine decision unaddressed at exactly the point a learner needs help. You will end up making this choice from forum posts and GitHub activity, which is not ideal. Debugging is the third issue and it is inherent rather than a documentation failure. When something goes wrong inside a jitted function, the error messages are improving but still rough, and the tools available are limited compared to stepping through eager PyTorch.
The docs cover the debugging utilities honestly and those utilities help. They do not make it pleasant. The last thing is scope. There is very little here on the actual practice of training models, on data loading, on checkpointing, on the workflow of an experiment.
Those are the daily concerns of someone doing machine learning and the docs treat them as out of scope, which is defensible for a numerical computing library and frustrating for a learner. On who should bother at all. If you are doing standard supervised learning on standard architectures, PyTorch is the pragmatic choice and JAX buys you very little. JAX earns its keep when you need transformations composed in unusual ways, when you are doing research that involves per sample gradients or higher order derivatives or aggressive vectorisation, or when you are working at a scale where the compiler's optimisations matter.
Reinforcement learning, scientific computing and large scale training are where it wins. My 4.0 is for outstanding writing constrained by format. As documentation it is close to the best in the field. As a course it is missing everything that makes a course work, and the score reflects both halves of that.