Skip to content

Functions

danling.tensor.functions

TorchFuncRegistry

Bases: Registry

Registry for extending PyTorch functions to work with custom tensor types like NestedTensor.

TorchFuncRegistry provides a clean interface for implementing PyTorch function overrides for custom tensor types such as NestedTensor. It’s used internally by NestedTensor to register implementations for various torch functions like torch.cat, torch.mean, torch.stack, etc.

This mechanism enables NestedTensor to behave like a regular torch.Tensor when used with standard PyTorch functions by providing custom implementations that understand the NestedTensor structure.

Usage:

Python
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
# Create a registry
registry = TorchFuncRegistry("my_tensor_registry")

# Register an implementation for torch.mean
@registry.implement(torch.mean)
def mean_implementation(input, dim=None, keepdim=False, **kwargs):
    # Custom implementation for your tensor type
    pass

# The registry can be used to look up the implementation
registry[torch.mean]  # Returns mean_implementation

Source code in danling/tensor/functions.py
Python
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
class TorchFuncRegistry(Registry):  # pylint: disable=too-few-public-methods
    """
    Registry for extending PyTorch functions to work with custom tensor types like NestedTensor.

    `TorchFuncRegistry` provides a clean interface for implementing PyTorch function
    overrides for custom tensor types such as NestedTensor. It's used internally by
    NestedTensor to register implementations for various torch functions like
    torch.cat, torch.mean, torch.stack, etc.

    This mechanism enables NestedTensor to behave like a regular torch.Tensor
    when used with standard PyTorch functions by providing custom implementations
    that understand the NestedTensor structure.

    Usage:
    ```python
    # Create a registry
    registry = TorchFuncRegistry("my_tensor_registry")

    # Register an implementation for torch.mean
    @registry.implement(torch.mean)
    def mean_implementation(input, dim=None, keepdim=False, **kwargs):
        # Custom implementation for your tensor type
        pass

    # The registry can be used to look up the implementation
    registry[torch.mean]  # Returns mean_implementation
    ```
    """

    def implement(self, torch_function: Callable) -> Callable:
        r"""
        Register a custom implementation for a PyTorch function.

        Use this decorator to provide implementations for PyTorch functions
        that will work with custom tensor types like NestedTensor. This is
        the key mechanism that allows NestedTensor to integrate seamlessly
        with the PyTorch ecosystem.

        Args:
            torch_function: The original PyTorch function to override (e.g., torch.mean, torch.cat)

        Returns:
            Callable: A decorator function that registers the implementation

        Raises:
            ValueError: If the function is already registered and override=False

        Examples:
            >>> import torch
            >>> registry = TorchFuncRegistry("test")
            >>> @registry.implement(torch.mean)
            ... def mean(input):
            ...     return input.mean()
            >>> registry[torch.mean]  # doctest: +ELLIPSIS
            <function mean at ...>

        Note:
            This is primarily used internally by NestedTensor.__torch_function__
            to provide implementations for various PyTorch functions. You can
            use the same mechanism to extend NestedTensor with additional
            function implementations.
        """

        if torch_function in self and not self.override:
            raise ValueError(f"Torch function {torch_function.__name__} already registered.")

        @wraps(self.register)
        def register(function):
            self.set(torch_function, function)
            return function

        return register

implement

Python
implement(torch_function: Callable) -> Callable

Register a custom implementation for a PyTorch function.

Use this decorator to provide implementations for PyTorch functions that will work with custom tensor types like NestedTensor. This is the key mechanism that allows NestedTensor to integrate seamlessly with the PyTorch ecosystem.

Parameters:

Name Type Description Default
torch_function
Callable

The original PyTorch function to override (e.g., torch.mean, torch.cat)

 required

Returns:

Name Type Description
Callable Callable

A decorator function that registers the implementation

Raises:

Type Description
ValueError

If the function is already registered and override=False

Examples:

Python Console Session
1
2
3
4
5
6
7
>>> import torch
>>> registry = TorchFuncRegistry("test")
>>> @registry.implement(torch.mean)
... def mean(input):
...     return input.mean()
>>> registry[torch.mean]
<function mean at ...>
Note

This is primarily used internally by NestedTensor.torch_function to provide implementations for various PyTorch functions. You can use the same mechanism to extend NestedTensor with additional function implementations.

Source code in danling/tensor/functions.py
Python
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
def implement(self, torch_function: Callable) -> Callable:
    r"""
    Register a custom implementation for a PyTorch function.

    Use this decorator to provide implementations for PyTorch functions
    that will work with custom tensor types like NestedTensor. This is
    the key mechanism that allows NestedTensor to integrate seamlessly
    with the PyTorch ecosystem.

    Args:
        torch_function: The original PyTorch function to override (e.g., torch.mean, torch.cat)

    Returns:
        Callable: A decorator function that registers the implementation

    Raises:
        ValueError: If the function is already registered and override=False

    Examples:
        >>> import torch
        >>> registry = TorchFuncRegistry("test")
        >>> @registry.implement(torch.mean)
        ... def mean(input):
        ...     return input.mean()
        >>> registry[torch.mean]  # doctest: +ELLIPSIS
        <function mean at ...>

    Note:
        This is primarily used internally by NestedTensor.__torch_function__
        to provide implementations for various PyTorch functions. You can
        use the same mechanism to extend NestedTensor with additional
        function implementations.
    """

    if torch_function in self and not self.override:
        raise ValueError(f"Torch function {torch_function.__name__} already registered.")

    @wraps(self.register)
    def register(function):
        self.set(torch_function, function)
        return function

    return register

NestedTensorFuncWrapper

Function Wrapper to handle NestedTensor as input.

Source code in danling/tensor/functions.py
Python
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
class NestedTensorFuncWrapper:  # pylint: disable=R0903
    r"""
    Function Wrapper to handle NestedTensor as input.
    """

    __storage: Sequence[Callable] = []
    state: Mapping = {}

    def __init__(self, *callables: Iterable[Callable], state: Mapping | None = None) -> None:
        if len(callables) == 1 and isinstance(callables, Sequence):
            callables = callables[0]  # type: ignore
        self._storage = callables  # type: ignore
        if state is None:
            state = {}
        self.state = state
        self.device = self.state.get("device")

    @property
    def _storage(self):
        return self.__storage

    @_storage.setter
    def _storage(self, callables: Sequence):
        if not isinstance(callables, Sequence):
            raise ValueError(f"callables must be a Sequence, bug got {type(callables)}")
        if len(callables) == 0:
            raise ValueError("callables must be a non-empty Sequence.")
        if not callable(callables[0]):
            raise ValueError(f"callables must be a Sequence of Callable, bug got {type(callables[0])}")
        self.__storage = callables

    def __call__(self, *args, **kwargs) -> NestedTensor | Tensor | Sequence[Tensor]:
        from .nested_tensor import NestedTensor

        ret = [call(*args, **kwargs) for call in self._storage]
        elem = ret[0]
        if isinstance(elem, Tensor):
            try:
                return torch.stack(ret, dim=0)
            except (ValueError, RuntimeError):
                return NestedTensor(ret, **self.state)
        if elem.__hash__ is not None and len(set(ret)) == 1:
            return elem
        return ret