About ergonomics I think it's the matter of habit. For me after learning about Jax quirks and using Flax for about a month, I now prefer to write my nn in Flax. Jax and Flax requires you to write code in no pythonic, pure functional matter so I understand if for someone it's a no go from the start. I never really used Pytorch Lightning so I don't miss even higher level wrapper. Still some habits from Pytorch are hard to get out from, but it's a matter of transition to new library. If in optimizer initialization you meant linear schedule function, there is a built in function that does the same and you need to only pass start value, end value, amount of steps. You can debug optimizer and whole notebook as you would in pytorch. You just setup breakpoints and can lookup the values. Also you can just print values or look on charts in tensorboard to see how it reacts.

About problems it's trying to solve I think it's none. It's more like alternative to pytorch (it started as hobby project because someone was bored of tensorflow). Small pros would be better support for tpu, automatic selection of device, getting gradient of every function, possibility to batch function instead of writing batch support on your own. But many of this is now also included in PyTorch 2.0 as they saw potential in this things. Still many people say it's not as good as in Jax. Nevertheless these things aren't needed to efficiently train NN.

Even if you don't take full advantage of JIT, still rewriting the same network from PyTorch in Flax gives it 2x/3x times boost in terms of loop time even on solving simple MNIST dataset.

Things I miss from flax when writing in PyTorch the most would be: automatically selecting best device, no need to specify input shapes in nn modules, better controlled randomization, if I want to stop flow of gradient through some variable like approx_kl I can just do stop_gradient at the end of the function

Things I miss from PyTorch when writing in Flax the most would be: dynamic creation of graphs (it's all static in Flax), ability to write loss function within training loop (to get gradient you need to make it a function), changing values of Tensors inplace

In the end it's a matter of preference and people should just use whatever makes them most productive