chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,504 @@
|
||||
# Copyright (c) 2021 PaddlePaddle Authors. 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.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING, Any, Concatenate, TypeVar
|
||||
|
||||
import paddle
|
||||
from paddle.base import core
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Callable, Sequence
|
||||
|
||||
from paddle import Tensor
|
||||
|
||||
|
||||
__all__ = []
|
||||
|
||||
|
||||
_RetT = TypeVar('_RetT')
|
||||
|
||||
|
||||
class PyLayerContext:
|
||||
"""
|
||||
``PyLayerContext`` can assist the :ref:`api_paddle_autograd_PyLayer` in implementing certain functionalities.
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> import paddle
|
||||
>>> from paddle.autograd import PyLayer
|
||||
|
||||
>>> class cus_tanh(PyLayer):
|
||||
... @staticmethod
|
||||
... def forward(ctx, x):
|
||||
... # ctx is a object of PyLayerContext.
|
||||
... y = paddle.tanh(x)
|
||||
... ctx.save_for_backward(y)
|
||||
... return y
|
||||
...
|
||||
... @staticmethod
|
||||
... def backward(ctx, dy):
|
||||
... # ctx is a object of PyLayerContext.
|
||||
... (y,) = ctx.saved_tensor()
|
||||
... grad = dy * (1 - paddle.square(y))
|
||||
... return grad
|
||||
"""
|
||||
|
||||
container: tuple[Tensor, ...]
|
||||
not_inplace_tensors: tuple[Tensor, ...]
|
||||
non_differentiable: tuple[Tensor, ...]
|
||||
materialize_grads: bool
|
||||
grad_in_dtype_consistent: bool
|
||||
|
||||
def set_grad_in_dtype_consistent(self, flag: bool) -> None:
|
||||
"""
|
||||
Set whether to maintain gradient input dtype consistency between forward output and backward input.
|
||||
|
||||
Note:
|
||||
This API should be called only inside `forward`.
|
||||
By default, backward input gradients are automatically cast to match the dtype of forward outputs.
|
||||
Set this to `False` to disable automatic casting and maintain original gradient dtypes in backward.
|
||||
|
||||
Args:
|
||||
flag (bool): Whether to enable automatic dtype conversion in backward.
|
||||
- `True`: Cast backward input gradient to match forward output dtype (default behavior)
|
||||
- `False`: Preserve original dtype of backward input gradient
|
||||
|
||||
Returns:
|
||||
None
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> import paddle
|
||||
>>> from paddle.autograd import PyLayer
|
||||
>>> paddle.seed(2025)
|
||||
>>> class cus_tanh(PyLayer):
|
||||
... @staticmethod
|
||||
... def forward(ctx, x):
|
||||
... y = paddle.tanh(x)
|
||||
... # Pass tensors to backward.
|
||||
... ctx.save_for_backward(y)
|
||||
... # The gradient input in the backward process
|
||||
... # will not be automatically cast to the dtype of the forward output.
|
||||
... ctx.set_grad_in_dtype_consistent(False)
|
||||
... return y
|
||||
...
|
||||
... @staticmethod
|
||||
... def backward(ctx, dy):
|
||||
...
|
||||
... # Get the tensors passed by forward.
|
||||
... (y,) = ctx.saved_tensor()
|
||||
... grad = dy * (1 - paddle.square(y))
|
||||
... return grad
|
||||
>>> class cus_tanh_cast_grad(PyLayer):
|
||||
... @staticmethod
|
||||
... def forward(ctx, x):
|
||||
... y = paddle.tanh(x)
|
||||
... # Pass tensors to backward.
|
||||
... ctx.save_for_backward(y)
|
||||
... return y
|
||||
...
|
||||
... @staticmethod
|
||||
... def backward(ctx, dy):
|
||||
... # Get the tensors passed by forward.
|
||||
... (y,) = ctx.saved_tensor()
|
||||
... grad = dy * (1 - paddle.square(y))
|
||||
... # The gradient input in cus_tanh be cast to bfloat16 manually,
|
||||
... # and cus_tanh will not cast the gradient to the dtype of the forward output.
|
||||
... grad = paddle.cast(grad, paddle.float16)
|
||||
... return grad
|
||||
>>> x = paddle.randn([3, 3]).astype("float32")
|
||||
>>> x.stop_gradient = False
|
||||
>>> y = cus_tanh.apply(x)
|
||||
>>> z = cus_tanh_cast_grad.apply(y)
|
||||
>>> z.sum().backward()
|
||||
|
||||
"""
|
||||
self.grad_in_dtype_consistent = flag
|
||||
|
||||
def save_for_backward(self, *tensors: Tensor) -> None:
|
||||
"""
|
||||
Saves given tensors that backward need. Use ``saved_tensor`` in the `backward` to get the saved tensors.
|
||||
|
||||
Note:
|
||||
This API should be called at most once, and only inside `forward`.
|
||||
|
||||
Args:
|
||||
tensors(list of Tensors): Tensors to be stored.
|
||||
|
||||
Returns:
|
||||
None
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> import paddle
|
||||
>>> from paddle.autograd import PyLayer
|
||||
|
||||
>>> class cus_tanh(PyLayer):
|
||||
... @staticmethod
|
||||
... def forward(ctx, x):
|
||||
... # ctx is a context object that store some objects for backward.
|
||||
... y = paddle.tanh(x)
|
||||
... # Pass tensors to backward.
|
||||
... ctx.save_for_backward(y)
|
||||
... return y
|
||||
...
|
||||
... @staticmethod
|
||||
... def backward(ctx, dy):
|
||||
... # Get the tensors passed by forward.
|
||||
... (y,) = ctx.saved_tensor()
|
||||
... grad = dy * (1 - paddle.square(y))
|
||||
... return grad
|
||||
|
||||
"""
|
||||
self.container = tensors
|
||||
|
||||
def saved_tensor(self) -> tuple[Tensor, ...]:
|
||||
"""
|
||||
Get the tensors stored by ``save_for_backward``.
|
||||
|
||||
Returns:
|
||||
list of Tensors or None: If context contains tensors stored by `save_for_backward`,
|
||||
then return these tensors, otherwise return None.
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> import paddle
|
||||
>>> from paddle.autograd import PyLayer
|
||||
|
||||
>>> class cus_tanh(PyLayer):
|
||||
... @staticmethod
|
||||
... def forward(ctx, x):
|
||||
... # ctx is a context object that store some objects for backward.
|
||||
... y = paddle.tanh(x)
|
||||
... # Pass tensors to backward.
|
||||
... ctx.save_for_backward(y)
|
||||
... return y
|
||||
...
|
||||
... @staticmethod
|
||||
... def backward(ctx, dy):
|
||||
... # Get the tensors passed by forward.
|
||||
... (y,) = ctx.saved_tensor()
|
||||
... grad = dy * (1 - paddle.square(y))
|
||||
... return grad
|
||||
"""
|
||||
return self.container
|
||||
|
||||
@property
|
||||
def saved_tensors(self):
|
||||
"""
|
||||
Get the tensors stored by ``save_for_backward``. This attribute is an alias for the method ``saved_tensor()``.
|
||||
|
||||
Returns:
|
||||
list of Tensors or None: If context contains tensors stored by `save_for_backward`,
|
||||
then return these tensors, otherwise return None.
|
||||
"""
|
||||
return self.saved_tensor()
|
||||
|
||||
def mark_not_inplace(self, *args: Tensor) -> None:
|
||||
"""
|
||||
Marks inputs as not inplace.
|
||||
This should be called at most once, only from inside the `forward` method,
|
||||
and all arguments should be Tensor inputs.
|
||||
|
||||
If the Tensor returned by `forward` method is the same as the Tensor input of forward,
|
||||
and this Tensor is marked as not_inplace, then Paddle will help the user create a new Tensor as output.
|
||||
Thereby preventing the auto grad information of the input Tensor from being overwritten.
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> import paddle
|
||||
|
||||
>>> class Exp(paddle.autograd.PyLayer):
|
||||
... @staticmethod
|
||||
... def forward(ctx, x):
|
||||
... ctx.mark_not_inplace(x)
|
||||
... return x
|
||||
...
|
||||
... @staticmethod
|
||||
... def backward(ctx, grad_output):
|
||||
... out = grad_output.exp()
|
||||
... return out
|
||||
|
||||
>>> paddle.seed(2023)
|
||||
>>> x = paddle.randn((1, 1))
|
||||
>>> x.stop_gradient = False
|
||||
>>> attn_layers = []
|
||||
>>> for idx in range(0, 2):
|
||||
... attn_layers.append(Exp())
|
||||
|
||||
>>> for step in range(0, 2):
|
||||
... a = x
|
||||
... for j in range(0, 2):
|
||||
... a = attn_layers[j].apply(x)
|
||||
... a.backward()
|
||||
"""
|
||||
self.not_inplace_tensors = args
|
||||
|
||||
def mark_non_differentiable(self, *args: Tensor) -> None:
|
||||
"""
|
||||
Marks outputs as non-differentiable.
|
||||
This should be called at most once, only from inside the `forward` method,
|
||||
and all arguments should be tensor outputs.
|
||||
|
||||
This will mark outputs as not requiring gradients, increasing the
|
||||
efficiency of backward computation. You still need to accept a gradient
|
||||
for each output in `backward`, but it's always going to
|
||||
be a zero tensor with the same shape as the shape of a corresponding
|
||||
output.
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> import paddle
|
||||
>>> from paddle.autograd import PyLayer
|
||||
>>> import numpy as np
|
||||
|
||||
>>> class Tanh(PyLayer):
|
||||
... @staticmethod
|
||||
... def forward(ctx, x):
|
||||
... a = x + x
|
||||
... b = x + x + x
|
||||
... ctx.mark_non_differentiable(a)
|
||||
... return a, b
|
||||
...
|
||||
... @staticmethod
|
||||
... def backward(ctx, grad_a, grad_b):
|
||||
... assert np.equal(grad_a.numpy(), paddle.zeros([1]).numpy())
|
||||
... assert np.equal(grad_b.numpy(), paddle.ones([1], dtype="float64").numpy())
|
||||
... return grad_b
|
||||
|
||||
>>> x = paddle.ones([1], dtype="float64")
|
||||
>>> x.stop_gradient = False
|
||||
>>> a, b = Tanh.apply(x)
|
||||
>>> b.sum().backward()
|
||||
"""
|
||||
self.non_differentiable = args
|
||||
|
||||
def set_materialize_grads(self, value: bool) -> None:
|
||||
"""
|
||||
Sets whether to materialize output grad tensors. Default is True.
|
||||
|
||||
This should be called only from inside the `forward` method.
|
||||
|
||||
If True, undefined output grad tensors will be expanded to tensors full
|
||||
of zeros prior to calling the `backward` method.
|
||||
|
||||
If False, undefined output grad tensors will be None.
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> import paddle
|
||||
>>> from paddle.autograd import PyLayer
|
||||
>>> import numpy as np
|
||||
|
||||
>>> class Tanh(PyLayer):
|
||||
... @staticmethod
|
||||
... def forward(ctx, x):
|
||||
... return x + x + x, x + x
|
||||
...
|
||||
... @staticmethod
|
||||
... def backward(ctx, grad, grad2):
|
||||
... assert np.equal(grad2.numpy(), paddle.zeros([1]).numpy())
|
||||
... return grad
|
||||
|
||||
>>> class Tanh2(PyLayer):
|
||||
... @staticmethod
|
||||
... def forward(ctx, x):
|
||||
... ctx.set_materialize_grads(False)
|
||||
... return x + x + x, x + x
|
||||
...
|
||||
... @staticmethod
|
||||
... def backward(ctx, grad, grad2):
|
||||
... assert grad2 == None
|
||||
... return grad
|
||||
|
||||
>>> x = paddle.ones([1], dtype="float64")
|
||||
>>> x.stop_gradient = False
|
||||
>>> Tanh.apply(x)[0].backward()
|
||||
|
||||
>>> x2 = paddle.ones([1], dtype="float64")
|
||||
>>> x2.stop_gradient = False
|
||||
>>> Tanh2.apply(x2)[0].backward()
|
||||
"""
|
||||
self.materialize_grads = value
|
||||
|
||||
|
||||
class PyLayerBackward(core.eager.PyLayer, PyLayerContext):
|
||||
def backward(self, *args):
|
||||
return self._forward_cls.backward(self, *args)
|
||||
|
||||
|
||||
class PyLayerMeta(type):
|
||||
def __init__(cls, name, bases, attrs):
|
||||
cls._backward_function = type(
|
||||
name + '_backward', (PyLayerBackward,), {"_forward_cls": cls}
|
||||
)
|
||||
|
||||
super().__init__(name, bases, attrs)
|
||||
|
||||
|
||||
class PyLayer(core.eager.PyLayer, PyLayerContext, metaclass=PyLayerMeta):
|
||||
"""
|
||||
Paddle implements Python custom operators on the PaddlePaddle framework by creating a subclass of
|
||||
``PyLayer``, which must comply with the following rules:
|
||||
|
||||
1. The subclass must contain static ``forward`` and ``backward`` functions, with the first argument being
|
||||
:ref:`api_paddle_autograd_PyLayerContext`. If a returned value in ``backward`` corresponds to a ``Tensor`` that
|
||||
requires gradients in ``forward``, the returned value must be a ``Tensor``.
|
||||
|
||||
2. Except for the first argument, other arguments of ``backward`` are gradients of the output ``Tensors``
|
||||
of ``forward``. Therefore, the number of input ``Tensor`` in ``backward`` must be the same as the number
|
||||
of output ``Tensor`` in ``forward``. If you need to use input ``Tensor`` from ``forward`` in ``backward``,
|
||||
you can save these ``Tensors`` by inputting them into :ref:`api_paddle_autograd_PyLayerContext`'s
|
||||
``save_for_backward`` method and use them in ``backward`` later.
|
||||
|
||||
3. The output of ``backward`` can be ``Tensor`` or ``list/tuple(Tensor)``, which are gradients of the
|
||||
output ``Tensor`` of ``forward``. Therefore, the number of output ``Tensor`` in ``backward`` is the same
|
||||
as the number of input ``Tensor`` in ``forward``.
|
||||
|
||||
After building the custom operator, apply it by running the ``apply`` method.
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> import paddle
|
||||
>>> from paddle.autograd import PyLayer
|
||||
|
||||
>>> class cus_tanh(PyLayer):
|
||||
... @staticmethod
|
||||
... def forward(ctx, x):
|
||||
... y = paddle.tanh(x)
|
||||
... # Pass tensors to backward.
|
||||
... ctx.save_for_backward(y)
|
||||
... return y
|
||||
...
|
||||
... @staticmethod
|
||||
... def backward(ctx, dy):
|
||||
... # Get the tensors passed by forward.
|
||||
... (y,) = ctx.saved_tensor()
|
||||
... grad = dy * (1 - paddle.square(y))
|
||||
... return grad
|
||||
|
||||
>>> paddle.seed(2023)
|
||||
>>> data = paddle.randn([2, 3], dtype="float64")
|
||||
>>> data.stop_gradient = False
|
||||
>>> z = cus_tanh.apply(data)
|
||||
>>> z.mean().backward()
|
||||
|
||||
>>> print(data.grad)
|
||||
Tensor(shape=[2, 3], dtype=float64, place=Place(cpu), stop_gradient=True,
|
||||
[[0.16604150, 0.05858341, 0.14051214],
|
||||
[0.15677770, 0.01564609, 0.02991660]])
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def forward(
|
||||
ctx: PyLayerContext, *args: Any, **kwargs: Any
|
||||
) -> Tensor | Sequence[Tensor]:
|
||||
"""
|
||||
It is to be overloaded by subclasses. It must accept a object of :ref:`api_paddle_autograd_PyLayerContext` as
|
||||
the first argument, followed by any number of arguments (tensors or other types).
|
||||
`None` can not be included in the returned result.
|
||||
|
||||
Args:
|
||||
*args(tuple): input of PyLayer.
|
||||
**kwargs(dict): input of PyLayer.
|
||||
|
||||
Returns:
|
||||
tensors or other types : output of PyLayer.
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> import paddle
|
||||
>>> from paddle.autograd import PyLayer
|
||||
|
||||
>>> class cus_tanh(PyLayer):
|
||||
... @staticmethod
|
||||
... def forward(ctx, x):
|
||||
... y = paddle.tanh(x)
|
||||
... # Pass tensors to backward.
|
||||
... ctx.save_for_backward(y)
|
||||
... return y
|
||||
...
|
||||
... @staticmethod
|
||||
... def backward(ctx, dy):
|
||||
... # Get the tensors passed by forward.
|
||||
... (y,) = ctx.saved_tensor()
|
||||
... grad = dy * (1 - paddle.square(y))
|
||||
... return grad
|
||||
"""
|
||||
raise NotImplementedError(
|
||||
"You must implement the forward function for PyLayer."
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx: PyLayerContext, *args: Any) -> Tensor | Sequence[Tensor]:
|
||||
"""
|
||||
This is a function to calculate the gradient. It is to be overloaded by subclasses.
|
||||
It must accept a object of :ref:`api_paddle_autograd_PyLayerContext` as the first
|
||||
argument, and the rest arguments are the gradient of forward's output tensors.
|
||||
Output tensors of backward are the gradient of forward's input tensors.
|
||||
|
||||
Args:
|
||||
*args(tuple): The gradient of forward's output tensor(s).
|
||||
**kwargs(dict): The gradient of forward's output tensor(s).
|
||||
|
||||
Returns:
|
||||
Tensor or list of Tensors: The gradient of forward's input tensor(s).
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> import paddle
|
||||
>>> from paddle.autograd import PyLayer
|
||||
|
||||
>>> class cus_tanh(PyLayer):
|
||||
... @staticmethod
|
||||
... def forward(ctx, x):
|
||||
... y = paddle.tanh(x)
|
||||
... # Pass tensors to backward.
|
||||
... ctx.save_for_backward(y)
|
||||
... return y
|
||||
...
|
||||
... @staticmethod
|
||||
... def backward(ctx, dy):
|
||||
... # Get the tensors passed by forward.
|
||||
... (y,) = ctx.saved_tensor()
|
||||
... grad = dy * (1 - paddle.square(y))
|
||||
... return grad
|
||||
"""
|
||||
|
||||
raise NotImplementedError(
|
||||
"You must implement the backward function for PyLayer."
|
||||
)
|
||||
|
||||
|
||||
def once_differentiable(
|
||||
backward: Callable[Concatenate[PyLayerContext, ...], _RetT],
|
||||
) -> Callable[Concatenate[PyLayerContext, ...], _RetT]:
|
||||
def wrapper(ctx: PyLayerContext, *args: Any) -> _RetT:
|
||||
with paddle.base.dygraph.no_grad():
|
||||
outputs = backward(ctx, *args)
|
||||
return outputs
|
||||
|
||||
return wrapper
|
||||
Reference in New Issue
Block a user