int64_t ret_index =
          mlir::cast<IntegerAttr>(
              execute.getDeviceVarUpdatesIndices().getValue()[entry.index()])
              .getInt();
      arg_metadata->set_retval_index_for_sharding(ret_index);
    }
  }
  // Update the metadata of the compile op.
  compile.setAttr("metadata", StringAttr::get(compile.getContext(),