Comment by one-more-minute

Right. It's important to remember that tools like JAX and PyTorch have total control over the numerical libraries they are differentiating, and have freedom to impose whatever semantics, rules and restrictions are convenient (immutability and referential transparency in JAX, for example). Seemingly small decisions in an existing language and library can have a big impact on the feasibility and practicality of AD.