About automatic distributed parallelism in JAX

[south-ocean]

I see that Alpa can implement automatic distributed programs for Jax programs through decorator @parallel, which can implement Data Parallel, Tensor Parallel, and Pipeline Parallel. And the core idea of Alpa has been incorporated into XLA for a long time, so there is a decorator method like Alpa's Parallel in JAX to achieve automatic distributed parallelism? I've only seen pmap, vmap, etc., both of which are data parallel correlation.
Whether there is an automatic distributed data connection, or the JAX design concept is to let the developer choose the appropriate parallelization strategy according to the specific needs and scenarios, rather than the framework to provide automatic distributed parallelism. Thanks.