add values_as_arrays arg

Author: nv-rliu

nx_cugraph/classes/graph.py

@@ -1171,17 +1171,27 @@ def _nodearray_to_set(self, node_ids: cp.ndarray[IndexValue]) -> set[NodeKey]:
     def _nodearray_to_dict(
         self,
         values: cp.ndarray[NodeValue],
+        values_as_arrays: bool = False,
     ) -> dict[NodeKey, NodeValue]:
-        # values_as_arrays: bool | None = None,
-        it = enumerate(values.tolist())
+        if values_as_arrays:
+            it = enumerate(cp.asnumpy(values))
+        else:
+            it = enumerate(values.tolist())
         if (id_to_key := self.id_to_key) is not None:
             return {id_to_key[key]: val for key, val in it}
         return dict(it)
 
     def _nodearrays_to_dict(
-        self, node_ids: cp.ndarray[IndexValue], values: any_ndarray[NodeValue]
+        self,
+        node_ids: cp.ndarray[IndexValue],
+        values: any_ndarray[NodeValue],
+        values_as_arrays: bool = False,
     ) -> dict[NodeKey, NodeValue]:
-        it = zip(node_ids.tolist(), values.tolist())
+        if values_as_arrays:
+            vals = cp.asnumpy(values)
+        else:
+            vals = values.tolist()
+        it = zip(node_ids.tolist(), vals)
         if (id_to_key := self.id_to_key) is not None:
             return {id_to_key[key]: val for key, val in it}
         return dict(it)