It's mostly an ecosystem thing, being unable to use existing methods.
In my experience, research goes something like
1. Milestone paper introducing novel method is published with green-field implementation
2. Bunch of papers extend milestone paper with brown-field implementation
3. Goto 1
Most things in 1 are written in PyTorch, meaning 2 also has to be in PyTorch.
I know this isn't JAX's fault, but I don't think JAX's philosophy to stay unopinionated and low-level is helping. Seems like the community agreeing on a single set of DL libraries around JAX will help it gain some momentum.
1. Milestone paper introducing novel method is published with green-field implementation
2. Bunch of papers extend milestone paper with brown-field implementation
3. Goto 1
Most things in 1 are written in PyTorch, meaning 2 also has to be in PyTorch. I know this isn't JAX's fault, but I don't think JAX's philosophy to stay unopinionated and low-level is helping. Seems like the community agreeing on a single set of DL libraries around JAX will help it gain some momentum.