Equinox has great idioms — really pioneered the Pytree perspective. Penzai is also great.
JAX feels close to achieving a sort of high-level GPGPU ecosystem. Super fledgling — but I keep finding more little libraries that build on JAX (and can be compositionally used with other JAX libraries because of it).
Only problem is that lots of compositional usage leads to big code and therefore big compile times for XLA.