How to ppermute a dynamic shape tensor when using shard_map?

[MoFHeka]

For example

input = [0,1,2,3,4,5,6,7]

The input is regarded as one training sample and segmented in the batch dimension with shard_map. The 'i' is the dimension name.

device 0:
input = [0,1,2,3]

device 1:
input = [4,5,6,7]

And I need to ppermute the input between different devices.

device 0:
lax.ppermute(input[0:1], 'i', [(0,1),(1,0)])

device 1:
lax.ppermute(input[1:3], 'i', [(0,1),(1,0)])

Where jax.disable_jit() is useless.

---

Dispatching parallel data input dynamically is an important way to handle Embedding model parallelism in the recommender model. Please, tell me how to support this feature.

Here is a snapshot from Meta paper :

Please:

• Check for duplicate requests.
• Describe your goal, and if possible provide a code snippet with a motivating example.