Sure, but OP's point is that TensorFlow's built-in libraries suck and you're forced to use them. They are confusing, change all the time, etc etc. Since JAX's

Sure, but OP's point is that TensorFlow's built-in libraries suck and you're forced to use them. They are confusing, change all the time, etc etc. Since JAX's API is lower-level, you can choose a neural network API that doesn't have the problems of TF's API. Or worst case, you make it yourself; you can see that the higher-level APIs provided by JAX are all very simple.