Repository: Project-MONAI/MONAI
File: monai/data/utils.py (Python)

Code before:

def pickle_operations(data, key=PICKLE_KEY_SUFFIX, is_encode: bool = True):
    """
    Applied_operations are dictionaries with varying sizes, this method converts them to bytes so that we can (de-)collate.

    Args:
        data: a list or dictionary with substructures to be pickled/unpickled.
        key: the key suffix for the target substructures, defaults to "_transforms" (`data.utils.PICKLE_KEY_SUFFIX`).
        is_encode: whether it's encoding using pickle.dumps (True) or decoding using pickle.loads (False).
    """
    if isinstance(data, Mapping):
        data = dict(data)
        for k in data:
            if f"{k}".endswith(key):
                if is_encode and not isinstance(data[k], bytes):
                    data[k] = pickle.dumps(data[k], 0)
                if not is_encode and isinstance(data[k], bytes):
                    data[k] = pickle.loads(data[k])
        return {k: pickle_operations(v, key=key, is_encode=is_encode) for k, v in data.items()}
    elif isinstance(data, (list, tuple)):
        return [pickle_operations(item, key=key, is_encode=is_encode) for item in data]
    return data


def collate_meta_tensor_fn(batch, *, collate_fn_map=None):
    """collate a sequence of meta tensor into a single batched metatensor"""
    collate_fn = collate_tensor_fn if pytorch_after(1, 13) else default_collate
    collated = collate_fn(batch)  # type: ignore
    meta_dicts = [i.meta or TraceKeys.NONE for i in batch]
    common_ = set.intersection(*[set(d.keys()) for d in meta_dicts if isinstance(d, dict)])
    if common_:
        meta_dicts = [{k: d[k] for k in common_} if isinstance(d, dict) else TraceKeys.NONE for d in meta_dicts]
    collated.meta = default_collate(meta_dicts)
    collated.applied_operations = [i.applied_operations or TraceKeys.NONE for i in batch]
    collated.is_batch = True
    return collated


def collate_meta_tensor(batch):
    """collate a sequence of meta tensor sequences/dictionaries into
    a single batched metatensor or a dictionary of batched metatensor"""
    if not isinstance(batch, Sequence):
        raise NotImplementedError()
    elem_0 = first(batch)
    if isinstance(elem_0, MetaObj):
        return collate_meta_tensor_fn(batch)
    if isinstance(elem_0, Mapping):
        return {k: collate_meta_tensor([d[k] for d in batch]) for k in elem_0}
    if isinstance(elem_0, (tuple, list)):
        return [collate_meta_tensor([d[i] for d in batch]) for i in range(len(elem_0))]


Review comment (performance):
```suggestion
    """
    Collate a sequence of meta tensor into a single batched metatensor. This is called by `collage_meta_tensor` and so should not be used as a collate function directly in dataloaders.
    """
```

I think we should say a little more to make it clear what this function is and to use `collage_meta_tensor` instead.

Code after:


def pickle_operations(data, key=PICKLE_KEY_SUFFIX, is_encode: bool = True):
    """
    Applied_operations are dictionaries with varying sizes, this method converts them to bytes so that we can (de-)collate.

    Args:
        data: a list or dictionary with substructures to be pickled/unpickled.
        key: the key suffix for the target substructures, defaults to "_transforms" (`data.utils.PICKLE_KEY_SUFFIX`).
        is_encode: whether it's encoding using pickle.dumps (True) or decoding using pickle.loads (False).
    """
    if isinstance(data, Mapping):
        data = dict(data)
        for k in data:
            if f"{k}".endswith(key):
                if is_encode and not isinstance(data[k], bytes):
                    data[k] = pickle.dumps(data[k], 0)
                if not is_encode and isinstance(data[k], bytes):
                    data[k] = pickle.loads(data[k])
        return {k: pickle_operations(v, key=key, is_encode=is_encode) for k, v in data.items()}
    elif isinstance(data, (list, tuple)):
        return [pickle_operations(item, key=key, is_encode=is_encode) for item in data]
    return data


def collate_meta_tensor_fn(batch, *, collate_fn_map=None):
    """
    Collate a sequence of meta tensor into a single batched metatensor. This is called by `collage_meta_tensor`
    and so should not be used as a collate function directly in dataloaders.
    """
    collate_fn = collate_tensor_fn if pytorch_after(1, 13) else default_collate
    collated = collate_fn(batch)  # type: ignore
    meta_dicts = [i.meta or TraceKeys.NONE for i in batch]
    common_ = set.intersection(*[set(d.keys()) for d in meta_dicts if isinstance(d, dict)])
    if common_:
        meta_dicts = [{k: d[k] for k in common_} if isinstance(d, dict) else TraceKeys.NONE for d in meta_dicts]
    collated.meta = default_collate(meta_dicts)
    collated.applied_operations = [i.applied_operations or TraceKeys.NONE for i in batch]
    collated.is_batch = True
    return collated


def collate_meta_tensor(batch):
    """collate a sequence of meta tensor sequences/dictionaries into
    a single batched metatensor or a dictionary of batched metatensor"""
    if not isinstance(batch, Sequence):
        raise NotImplementedError()
    elem_0 = first(batch)
    if isinstance(elem_0, MetaObj):
        return collate_meta_tensor_fn(batch)
    if isinstance(elem_0, Mapping):