Wanna know how to extract the intermediate tensors' partition strategies?

[LeiWang1999]

Hey everyone! I'm working on extracting tensor parallel or shard parallel strategies from alpa's auto-parallelization feature.

Using the quick start guide from alpa.ai as a reference, I've managed to retrieve the input, output tensor, and constant data of Dense Layers from actual_state and batch.

But I'm a bit stuck on how to pull out the parallel strategies for intermediate tensors, especially the ones between ‘Dense_0’ and 'Dense_1'. Any insights or advice on this would be greatly appreciated!