Use the rollout engine construct mesh.

[wang2yn84]

We see quite a few breakages around how mesh is constructed differently in rollout engine. This PR fixes this issue by using the mesh constructed by rollout engine to get rid of this discrepancies completely.

Reference

Colab Notebook

Checklist

• I have added all the necessary unit tests for my change.
• I have verified that my change does not break existing code and all unit tests pass.
• I have added all appropriate doc-strings/documentation.
• My PR is based on the latest changes of the main branch (if unsure, rebase the code).
• I have signed the Contributor License Agreement.
• I have followed Contribution Guidelines.

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!