// order operations on device. First op in the region should have a special
// argument token and then remaining operations should have node name of the
// previous communication ops.
LogicalResult SetTokenInputAttrs(ModuleOp module) {
  // Collect all the ops that needs to have token input names attributes. These
  // ops are communication ops and all their parent ops via nesting or function