Function tensorfun

cdslib/core/utils/tensor_utils.py:10–42

Apply `func(tensor)` to any tensors in the list or dict. Args: func: function that takes a tensor and outputs a tensor arr: list of tensor, dict of tensor Returns: list or dict containing the output tensors

(
    func: T.Callable[[torch.Tensor], torch.Tensor],
    arr: T.Union[torch.Tensor, T.List[torch.Tensor], T.Dict[T.Any, torch.Tensor]],
)

def tensorfun(
    func: T.Callable[[torch.Tensor], torch.Tensor],
    arr: T.Union[torch.Tensor, T.List[torch.Tensor], T.Dict[T.Any, torch.Tensor]],
) -> T.Union[torch.Tensor, T.List[torch.Tensor], T.Dict[T.Any, torch.Tensor], None]:
    """
    Apply `func(tensor)` to any tensors in the list or dict.

    Args:
        func:
            function that takes a tensor and outputs a tensor
        arr:
            list of tensor, dict of tensor

    Returns:
        list or dict containing the output tensors
    """

    if arr is None:
        return None

    if isinstance(arr, torch.Tensor):
        return func(arr)

    if isinstance(arr, (list, tuple)):
        return [tensorfun(func, a) for a in arr]

    if isinstance(arr, dict):
        out_dict = dict()
        for key in arr:
            out_dict[key] = tensorfun(func, arr[key])
        return out_dict

    raise RuntimeError(f'Not supported type: {type(arr)}, {arr}')


def random_crop(