170 lines
6.3 KiB
Python
170 lines
6.3 KiB
Python
# Copyright 2022-2025 MetaOPT Team. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
# ==============================================================================
|
|
"""PyTree integration with :mod:`functools`."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import functools
|
|
from typing import TYPE_CHECKING, Any, Callable, ClassVar
|
|
from typing_extensions import Self # Python 3.11+
|
|
|
|
from optree import registry
|
|
from optree.accessor import GetAttrEntry
|
|
from optree.ops import tree_reduce as reduce
|
|
from optree.typing import CustomTreeNode, T
|
|
|
|
|
|
if TYPE_CHECKING:
|
|
from optree.accessor import PyTreeEntry
|
|
|
|
|
|
__all__ = [
|
|
'partial',
|
|
'reduce',
|
|
]
|
|
|
|
|
|
class _HashablePartialShim:
|
|
"""Object that delegates :meth:`__call__`, :meth:`__eq__`, and :meth:`__hash__` to another object."""
|
|
|
|
__slots__: ClassVar[tuple[str, ...]] = ('args', 'func', 'keywords', 'partial_func')
|
|
|
|
func: Callable[..., Any]
|
|
args: tuple[Any, ...]
|
|
keywords: dict[str, Any]
|
|
|
|
def __init__(self, partial_func: functools.partial, /) -> None:
|
|
self.partial_func: functools.partial = partial_func
|
|
|
|
def __call__(self, /, *args: Any, **kwargs: Any) -> Any:
|
|
return self.partial_func(*args, **kwargs)
|
|
|
|
def __eq__(self, other: object, /) -> bool:
|
|
if isinstance(other, _HashablePartialShim):
|
|
return self.partial_func == other.partial_func
|
|
return self.partial_func == other
|
|
|
|
def __hash__(self, /) -> int:
|
|
return hash(self.partial_func)
|
|
|
|
def __repr__(self, /) -> str:
|
|
return repr(self.partial_func)
|
|
|
|
|
|
# pylint: disable-next=protected-access
|
|
@registry.register_pytree_node_class(namespace=registry.__GLOBAL_NAMESPACE)
|
|
class partial( # noqa: N801 # pylint: disable=invalid-name,too-few-public-methods
|
|
functools.partial,
|
|
CustomTreeNode[T],
|
|
):
|
|
"""A version of :func:`functools.partial` that works in pytrees.
|
|
|
|
Use it for partial function evaluation in a way that is compatible with transformations,
|
|
e.g., ``partial(func, *args, **kwargs)``.
|
|
|
|
(You need to explicitly opt-in to this behavior because we did not want to give
|
|
:func:`functools.partial` different semantics than normal function closures.)
|
|
|
|
For example, here is a basic usage of :class:`partial` in a manner similar to
|
|
:func:`functools.partial`:
|
|
|
|
>>> import operator
|
|
>>> import torch
|
|
>>> add_one = partial(operator.add, torch.ones(()))
|
|
>>> add_one(torch.tensor([[1, 2], [3, 4]]))
|
|
tensor([[2., 3.],
|
|
[4., 5.]])
|
|
|
|
Pytree compatibility means that the resulting partial function can be passed as an argument
|
|
within tree-map functions, which is not possible with a standard :func:`functools.partial`
|
|
function:
|
|
|
|
>>> def call_func_on_cuda(f, *args, **kwargs):
|
|
... f, args, kwargs = tree_map(lambda t: t.cuda(), (f, args, kwargs))
|
|
... return f(*args, **kwargs)
|
|
...
|
|
>>> # doctest: +SKIP
|
|
>>> tree_map(lambda t: t.cuda(), add_one)
|
|
optree.functools.partial(<built-in function add>, tensor(1., device='cuda:0'))
|
|
>>> call_func_on_cuda(add_one, torch.tensor([[1, 2], [3, 4]]))
|
|
tensor([[2., 3.],
|
|
[4., 5.]], device='cuda:0')
|
|
|
|
Passing zero arguments to :class:`partial` effectively wraps the original function, making it a
|
|
valid argument in tree-map functions:
|
|
|
|
>>> # doctest: +SKIP
|
|
>>> call_func_on_cuda(partial(torch.add), torch.tensor(1), torch.tensor(2))
|
|
tensor(3, device='cuda:0')
|
|
|
|
Had we passed :func:`operator.add` to ``call_func_on_cuda`` directly, it would have resulted in
|
|
a :class:`TypeError` or :class:`AttributeError`.
|
|
"""
|
|
|
|
__slots__: ClassVar[tuple[()]] = ()
|
|
|
|
func: Callable[..., Any]
|
|
args: tuple[T, ...]
|
|
keywords: dict[str, T]
|
|
|
|
TREE_PATH_ENTRY_TYPE: ClassVar[type[PyTreeEntry]] = GetAttrEntry
|
|
|
|
def __new__(cls, func: Callable[..., Any], /, *args: T, **keywords: T) -> Self:
|
|
"""Create a new :class:`partial` instance."""
|
|
# In Python 3.10+, if func is itself a functools.partial instance, functools.partial.__new__
|
|
# would merge the arguments of this partial instance with the arguments of the func. We box
|
|
# func in a class that does not (yet) have a `func` attribute to defeat this optimization,
|
|
# since we care exactly which arguments are considered part of the pytree.
|
|
if isinstance(func, functools.partial):
|
|
original_func = func
|
|
func = _HashablePartialShim(original_func)
|
|
assert not hasattr(func, 'func'), 'shimmed function should not have a `func` attribute'
|
|
out = super().__new__(cls, func, *args, **keywords)
|
|
func.func = original_func.func
|
|
func.args = original_func.args
|
|
func.keywords = original_func.keywords
|
|
return out
|
|
|
|
return super().__new__(cls, func, *args, **keywords)
|
|
|
|
def __repr__(self, /) -> str:
|
|
"""Return a string representation of the :class:`partial` instance."""
|
|
args = [repr(self.func)]
|
|
args.extend(repr(x) for x in self.args)
|
|
args.extend(f'{k}={v!r}' for (k, v) in self.keywords.items())
|
|
return f'{self.__class__.__module__}.{self.__class__.__qualname__}({", ".join(args)})'
|
|
|
|
def tree_flatten( # type: ignore[override]
|
|
self,
|
|
/,
|
|
) -> tuple[
|
|
tuple[tuple[T, ...], dict[str, T]],
|
|
Callable[..., Any],
|
|
tuple[str, str],
|
|
]:
|
|
"""Flatten the :class:`partial` instance to children and metadata."""
|
|
return (self.args, self.keywords), self.func, ('args', 'keywords')
|
|
|
|
@classmethod
|
|
def tree_unflatten( # type: ignore[override]
|
|
cls,
|
|
metadata: Callable[..., Any],
|
|
children: tuple[tuple[T, ...], dict[str, T]],
|
|
/,
|
|
) -> Self:
|
|
"""Unflatten the children and metadata into a :class:`partial` instance."""
|
|
args, keywords = children
|
|
return cls(metadata, *args, **keywords)
|