Hey @Sai_Kiran, are you referring to Mesh-TensorFlow-like parallelism? We currently don’t have the API for it yet. But the split, scatter, parallel_apply, and gather can be done in the application side.
With parallel_apply, there won’t be a loop in the application code, but it internally uses a loop to launch multiple threads to process inputs in parallel. Here 2 is an example usage of parallel_apply.