I think that would be an interesting project, and something that is entirely feasible.
But personally, since developing a frontend requires significant effort, we have good support for PyTorch, and PyTorch is increasing adding JAX-inspired feature, I’d rather improve our support for PyTorch.