collate a sequence of meta tensor sequences/dictionaries into a single batched metatensor or a dictionary of batched metatensor
(batch)
| 436 | |
| 437 | |
| 438 | def collate_meta_tensor(batch): |
| 439 | """collate a sequence of meta tensor sequences/dictionaries into |
| 440 | a single batched metatensor or a dictionary of batched metatensor""" |
| 441 | if not isinstance(batch, Sequence): |
| 442 | raise NotImplementedError() |
| 443 | elem_0 = first(batch) |
| 444 | if isinstance(elem_0, MetaObj): |
| 445 | return collate_meta_tensor_fn(batch) |
| 446 | if isinstance(elem_0, Mapping): |
| 447 | return {k: collate_meta_tensor([d[k] for d in batch]) for k in elem_0} |
| 448 | if isinstance(elem_0, (tuple, list)): |
| 449 | return [collate_meta_tensor([d[i] for d in batch]) for i in range(len(elem_0))] |
| 450 | |
| 451 | # no more recursive search for MetaTensor |
| 452 | return default_collate(batch) |
| 453 | |
| 454 | |
| 455 | def list_data_collate(batch: Sequence): |
nothing calls this directly
no test coverage detected
searching dependent graphs…