Before anything else: check that you need JAX. It is the question this guide exists to make you ask, because the hiring consequences of the answer are larger than anything else on this page.

Across the UK careers pages we index, roughly 1,200 employers name PyTorch and about 200 name JAX. That is a six-fold difference in the pool you are recruiting from, and it does not reflect a six-fold difference in capability — it reflects where the industry standardised. If you are choosing a framework now, that ratio is a hiring cost you are signing up to. If the choice was made for you, the rest of this is about recruiting into it realistically.

The cases where JAX is genuinely the right answer

There are real ones, and they tend to look alike:

  • You are on TPUs. JAX is the first-class path there, and if your compute commitment is Google’s, the framework question is largely settled.
  • Your work is research-shaped — novel architectures, custom gradients, anything where you are differentiating through something unusual. The functional design makes transformations composable in a way that rewards this.
  • You need the compiler. Where XLA fusion gives you a material speed-up on your actual workload, that is a measured argument rather than a preference.
  • Your team already knows it, usually because somebody came from a lab that used it. Rewriting a working stack to widen a hiring pool is rarely worth it.

The case that does not hold up is choosing JAX because it is newer or faster in the abstract. For most production inference and fine-tuning work the ecosystem around PyTorch will get you there sooner with more people available to help.

What the framework demands that others do not

This matters for screening, because the gap is conceptual rather than syntactic. An engineer who is fluent in one imperative framework is not automatically productive here.

The specifics that catch people out: everything is functional and pure, so no in-place mutation and no hidden state. Randomness is explicit — you thread PRNG keys through your code by hand, and forgetting to split one produces correlated noise that silently ruins an experiment. Compilation happens under tracing, so Python control flow that depends on array values has to be rewritten, and a shape that changes triggers a recompile that can dominate your runtime. Parallelism is expressed by sharding arrays across devices rather than by wrapping a model.

Each of those is learnable in weeks by a strong engineer. None of them is guessable in an interview by somebody who has not met them.

Why almost nobody has five years of it, and what to do instead

The honest position is that a candidate pool of “JAX engineers” barely exists at the size most employers want. The deep experience sits in a small number of research organisations, and those people are neither numerous nor cheap.

What works better is to hire for the underlying capability and accept a short conversion. Look for strong numerical and array-programming instincts, evidence of caring about compilation and memory rather than only about model architecture, and ideally exposure to functional programming from any language at all. Somebody who has written serious NumPy, or who has used a compiler-oriented stack, converts quickly. An engineer whose entire experience is calling high-level training APIs converts slowly regardless of framework.

In practice we see teams succeed by hiring a very good machine learning engineer and budgeting a month, rather than waiting two quarters for somebody with the exact label.

Screening questions that work on a short-experience pool

Because you are largely assessing potential, the questions should test the concepts rather than the API.

“Why does this function recompile every time I call it?” Strong answers reach for changing input shapes or a traced value used in Python control flow. It is the single most common real-world JAX problem and the answer separates people who have run something from people who have read the documentation.

“How would you make sure two experiments are genuinely independent?” You want explicit key splitting and an understanding of why reusing a key silently correlates results. Somebody who says “set the seed” has not used the framework.

“Walk me through how you would shard this model across eight devices.” Listen for thinking about where the data lives and what has to be communicated, rather than for a remembered incantation.

For a candidate converting from elsewhere, replace the first two with the equivalent problems in their own stack. What you are testing is whether they debug the machine or the symptom.

What to pay, and why any figure here is soft

Indicatively, a strong engineer doing this work in Britain sits somewhere around £85,000 to £130,000 permanent, with contract day rates from roughly £600 upwards, higher where TPUs or large-scale training are involved.

Treat that as a very wide band, and here is the honest reason: with around 200 UK employers naming the framework at all, there is no sample from which a reliable benchmark could be built. Anybody quoting you a precise median for a JAX engineer has extrapolated it from adjacent roles. The practical consequence is that the market clears on the individual rather than on the title, and you should expect to negotiate against what the person could earn doing something else.

The route most teams actually end up taking

Two hires rather than one, staged. First, somebody senior who has genuinely used the framework in anger, even if that person is a contractor for three months. They set the patterns: the sharding strategy, the testing approach, the project structure that stops the codebase becoming un-compilable.

Then hire ordinary strong ML engineers into those patterns. That converts a scarce-hire problem into an abundant-hire problem, and it is considerably more robust than one specialist who becomes the only person who understands the training loop.

If you cannot get that first hire, the sober option is to reconsider the framework rather than to hire four people who will learn it together on your production system.