multi dim numpy array support

roll/datasets/collator.py

@@ -32,7 +32,8 @@ def collate_fn_to_dict_list(data_list: list[dict]) -> dict:
         tensors[key] = torch.cat(val, dim=0)
 
     for key, val in non_tensors.items():
-        non_tensors[key] = np.array(val, dtype=object)
+        non_tensors[key] = np.empty(len(val), dtype=object)
+        non_tensors[key][:] = val
 
     output = {}
     output.update(tensors)
@@ -214,5 +215,6 @@ def __call__(self, features: List[Dict[str, Any]]) -> Dict[str, Any]:
                 assert batch[key].shape[0] == batch["input_ids"].shape[0]
             else:
                 assert len(batch[key]) == batch["input_ids"].shape[0]
-                batch[key] = np.array(batch[key], dtype=object)
+                batch[key] = np.empty(len(batch[key]), dtype=object)
+                batch[key][:] = batch[key]
         return batch

roll/distributed/scheduler/protocol.py

@@ -109,7 +109,8 @@ def collate_fn(x: list["DataProtoItem"]):
     batch = torch.stack(batch).contiguous()
     non_tensor_batch = list_of_dict_to_dict_of_list(non_tensor_batch)
     for key, val in non_tensor_batch.items():
-        non_tensor_batch[key] = np.array(val, dtype=object)
+        non_tensor_batch[key] = np.empty(len(val), dtype=object)
+        non_tensor_batch[key][:] = val
     return DataProto(batch=batch, non_tensor_batch=non_tensor_batch, meta_info=meta_info)
 
 
@@ -286,7 +287,8 @@ def from_dict(cls, tensors: Dict[str, torch.Tensor], non_tensors=None, meta_info
                 ), f"Not all the tensor in tensors have the same batch size with batch_dims={num_batch_dims}. Got {pivot_key} has {batch_size}, {key} has {current_batch}"
 
         for key, val in non_tensors.items():
-            non_tensors[key] = np.array(val, dtype=object)
+            non_tensors[key] = np.empty(len(val), dtype=object)
+            non_tensors[key][:] = val
 
         tensor_dict = TensorDict(source=tensors, batch_size=batch_size)
         return cls(batch=tensor_dict, non_tensor_batch=non_tensors, meta_info=meta_info)