chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,44 @@
|
||||
# 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 ..base.dygraph.base import ( # noqa: F401
|
||||
enable_grad,
|
||||
grad,
|
||||
inference_mode,
|
||||
is_grad_enabled,
|
||||
no_grad_ as no_grad,
|
||||
set_grad_enabled,
|
||||
)
|
||||
from . import ( # noqa: F401
|
||||
backward_mode,
|
||||
function,
|
||||
grad_mode,
|
||||
ir_backward,
|
||||
)
|
||||
from .autograd import hessian, jacobian
|
||||
from .backward_mode import backward
|
||||
from .py_layer import PyLayer, PyLayerContext
|
||||
from .saved_tensors_hooks import saved_tensors_hooks
|
||||
|
||||
Function = PyLayer
|
||||
|
||||
__all__ = [
|
||||
'jacobian',
|
||||
'hessian',
|
||||
'backward',
|
||||
'PyLayer',
|
||||
'Function',
|
||||
'PyLayerContext',
|
||||
'saved_tensors_hooks',
|
||||
]
|
||||
@@ -0,0 +1,775 @@
|
||||
# Copyright (c) 2023 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 collections.abc import Sequence
|
||||
from typing import TYPE_CHECKING, overload
|
||||
|
||||
import paddle
|
||||
from paddle.base import framework
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from paddle import Tensor
|
||||
|
||||
|
||||
def as_tensors(xs):
|
||||
if isinstance(xs, framework.Variable):
|
||||
return xs
|
||||
elif isinstance(xs, Sequence):
|
||||
return tuple(xs)
|
||||
else:
|
||||
return xs
|
||||
|
||||
|
||||
class Jacobian:
|
||||
r"""Computes the Jacobian matrix of given xs and ys.
|
||||
|
||||
Once the Jacobian ``J`` is constructed, you can use a multidimensional index
|
||||
to retrieve the submatrix of ``J``, as same as slicing a Tensor. The
|
||||
submatrix is lazily evaluated along row axis, and will be cached once
|
||||
evaluated.
|
||||
|
||||
you can retrieve the submatrix by
|
||||
following methods:
|
||||
|
||||
* J[:], retrieving the full matrix.
|
||||
* J[:, :, j], retrieving the partial derivatives w.r.t. the j'th input
|
||||
variable.
|
||||
* J[:, i, :], retrieving the partial derivatives w.r.t. the i'th output
|
||||
variable.
|
||||
* J[:, i, j], retrieving the partial derivatives w.r.t. the i'th output
|
||||
variable and the j'th input variable.
|
||||
|
||||
Notes:
|
||||
|
||||
Ellipsis index is not supported currently.
|
||||
|
||||
Args:
|
||||
|
||||
ys (Tensor|Tuple[Tensor, ...]): The output derived from xs .
|
||||
xs (Tensor|Tuple[Tensor, ...]): The input tensor(s) .
|
||||
is_batched (bool): If true, the first axis is batch axis. Defaults to
|
||||
False.
|
||||
|
||||
Returns:
|
||||
|
||||
Jacobian (Object): A python object retains the Jacobian matrix.
|
||||
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
ys: Tensor,
|
||||
xs: Tensor,
|
||||
is_batched: bool = False,
|
||||
) -> None:
|
||||
if not is_batched:
|
||||
if not 0 <= len(xs.shape) <= 1:
|
||||
raise ValueError(
|
||||
f"xs.ndim should be 0 or 1 when is_batched=False"
|
||||
f" but got {len(xs.shape)}"
|
||||
)
|
||||
if not 0 <= len(ys.shape) <= 1:
|
||||
raise ValueError(
|
||||
f"ys.ndim should be 0 or 1 when is_batched=False"
|
||||
f" but got {len(ys.shape)}"
|
||||
)
|
||||
self._jacobian = _JacobianNoBatch(ys, xs)
|
||||
else:
|
||||
if not 1 <= len(ys.shape) <= 2:
|
||||
raise ValueError(
|
||||
f"ys.ndim should be 1 or 2 when is_batched=True"
|
||||
f" but got {len(ys.shape)}"
|
||||
)
|
||||
if not 1 <= len(xs.shape) <= 2:
|
||||
raise ValueError(
|
||||
f"xs.ndim should be 1 or 2 when is_batched=True"
|
||||
f" but got {len(xs.shape)}"
|
||||
)
|
||||
self._jacobian = _JacobianBatchFirst(ys, xs)
|
||||
|
||||
@property
|
||||
def shape(self) -> list[int]:
|
||||
"""The shape of flattened Jacobian matrix."""
|
||||
return self._jacobian.shape
|
||||
|
||||
def __getitem__(self, indexes):
|
||||
return self._jacobian[indexes]
|
||||
|
||||
def __getattr__(self, __name: str): # noqa: PYI063
|
||||
if __name == "shape":
|
||||
return getattr(self._jacobian, __name)
|
||||
if __name == "_evaluate_all":
|
||||
return getattr(self._jacobian, __name)
|
||||
return getattr(self._jacobian._evaluate_all(), __name)
|
||||
|
||||
def __add__(self, other):
|
||||
lhs = self._evaluate_all()
|
||||
rhs = other._evaluate_all() if isinstance(other, Jacobian) else other
|
||||
return lhs + rhs
|
||||
|
||||
def __sub__(self, other):
|
||||
lhs = self._evaluate_all()
|
||||
rhs = other._evaluate_all() if isinstance(other, Jacobian) else other
|
||||
return lhs - rhs
|
||||
|
||||
def __mul__(self, other):
|
||||
lhs = self._evaluate_all()
|
||||
rhs = other._evaluate_all() if isinstance(other, Jacobian) else other
|
||||
return lhs * rhs
|
||||
|
||||
def __div__(self, other):
|
||||
lhs = self._evaluate_all()
|
||||
rhs = other._evaluate_all() if isinstance(other, Jacobian) else other
|
||||
return lhs / rhs
|
||||
|
||||
def __truediv__(self, other):
|
||||
lhs = self._evaluate_all()
|
||||
rhs = other._evaluate_all() if isinstance(other, Jacobian) else other
|
||||
return lhs / rhs
|
||||
|
||||
def __pow__(self, other):
|
||||
lhs = self._evaluate_all()
|
||||
rhs = other._evaluate_all() if isinstance(other, Jacobian) else other
|
||||
return lhs**rhs
|
||||
|
||||
def __mod__(self, other):
|
||||
lhs = self._evaluate_all()
|
||||
rhs = other._evaluate_all() if isinstance(other, Jacobian) else other
|
||||
return lhs % rhs
|
||||
|
||||
def __floordiv__(self, other):
|
||||
lhs = self._evaluate_all()
|
||||
rhs = other._evaluate_all() if isinstance(other, Jacobian) else other
|
||||
return lhs // rhs
|
||||
|
||||
def __matmul__(self, other):
|
||||
lhs = self._evaluate_all()
|
||||
rhs = other._evaluate_all() if isinstance(other, Jacobian) else other
|
||||
return lhs @ rhs
|
||||
|
||||
def __eq__(self, other):
|
||||
lhs = self._evaluate_all()
|
||||
rhs = other._evaluate_all() if isinstance(other, Jacobian) else other
|
||||
return lhs == rhs
|
||||
|
||||
def __ne__(self, other):
|
||||
lhs = self._evaluate_all()
|
||||
rhs = other._evaluate_all() if isinstance(other, Jacobian) else other
|
||||
return lhs != rhs
|
||||
|
||||
def __lt__(self, other):
|
||||
lhs = self._evaluate_all()
|
||||
rhs = other._evaluate_all() if isinstance(other, Jacobian) else other
|
||||
return lhs < rhs
|
||||
|
||||
def __le__(self, other):
|
||||
lhs = self._evaluate_all()
|
||||
rhs = other._evaluate_all() if isinstance(other, Jacobian) else other
|
||||
return lhs <= rhs
|
||||
|
||||
def __gt__(self, other):
|
||||
lhs = self._evaluate_all()
|
||||
rhs = other._evaluate_all() if isinstance(other, Jacobian) else other
|
||||
return lhs > rhs
|
||||
|
||||
def __ge__(self, other):
|
||||
lhs = self._evaluate_all()
|
||||
rhs = other._evaluate_all() if isinstance(other, Jacobian) else other
|
||||
return lhs >= rhs
|
||||
|
||||
|
||||
class Hessian(Jacobian):
|
||||
pass
|
||||
|
||||
|
||||
class _Jacobian:
|
||||
"""The base class for computing Jacobian matrix.
|
||||
|
||||
``_Jacobian`` implements the core logic of multidimensional index and lazy
|
||||
evaluation for Jacobian matrix, subclass only need to overwrite following
|
||||
methods:
|
||||
|
||||
* ``_lazy_axis()``, return the axis along which will be lazy
|
||||
evaluating.
|
||||
* ``_flatten(xs)``, flattens the inputs ``xs``.
|
||||
* ``_evaluate(index)``, evaluates one slice along ``_lazy_axis`` .
|
||||
|
||||
Notes:
|
||||
|
||||
Because currently PaddlePaddle only support reverse differentiation by
|
||||
``paddle.grad``, so lazy evaluation is only supported along the row of
|
||||
Jacobian matrix, which means that slicing along row will get better
|
||||
performance.
|
||||
|
||||
"""
|
||||
|
||||
def __init__(self, ys, xs):
|
||||
self.original_xs_shape = xs.shape
|
||||
self.original_ys_shape = ys.shape
|
||||
self._xs = xs
|
||||
self._ys = ys
|
||||
if len(self._ys.shape) == 0 and not self.is_batched:
|
||||
self._ys = self._ys.reshape(
|
||||
[
|
||||
-1,
|
||||
]
|
||||
)
|
||||
if len(self._ys.shape) == 1 and self.is_batched:
|
||||
self._ys = self._ys.reshape([-1, 1])
|
||||
|
||||
self._flatten_xs = self._flatten(as_tensors(self._xs))
|
||||
self._flatten_ys = self._flatten(as_tensors(self._ys))
|
||||
self._cache = {}
|
||||
|
||||
@property
|
||||
def _lazy_axis(self):
|
||||
""" "The axis of lazily evaluated."""
|
||||
raise NotImplementedError
|
||||
|
||||
def _lazy_indexes(self, indexes):
|
||||
idx = indexes[self._lazy_axis]
|
||||
return (
|
||||
(idx,)
|
||||
if isinstance(idx, int)
|
||||
else tuple(range(idx.start, idx.stop, idx.step))
|
||||
)
|
||||
|
||||
def _flatten(self, xs):
|
||||
raise NotImplementedError
|
||||
|
||||
def _shifted_indexes(self, indexes, lazy_axis_size=0):
|
||||
idx = indexes[self._lazy_axis]
|
||||
shifted_lazy_axis_idx = (
|
||||
0 if isinstance(idx, int) else slice(0, lazy_axis_size, 1)
|
||||
)
|
||||
return (
|
||||
*indexes[: self._lazy_axis],
|
||||
shifted_lazy_axis_idx,
|
||||
*indexes[self._lazy_axis + 1 :],
|
||||
)
|
||||
|
||||
def __getitem__(self, indexes):
|
||||
if self.is_batched is False:
|
||||
if len(self.shape) == 0:
|
||||
# xs and ys are both 0-D tensor
|
||||
raise IndexError("0-D tensor can not be indexed.")
|
||||
elif len(self.shape) == 1:
|
||||
# either ys or xs is 0-D tensor
|
||||
indexes = (
|
||||
(0, indexes)
|
||||
if len(self.original_ys_shape) == 0
|
||||
else (indexes, 0)
|
||||
)
|
||||
else:
|
||||
if len(self.shape) == 1:
|
||||
# xs and ys are both 1-D tensor
|
||||
indexes = (indexes, 0, 0)
|
||||
elif len(self.shape) == 2:
|
||||
# either xs or ys is 1-D tensor
|
||||
if isinstance(indexes, slice):
|
||||
indexes = (indexes, slice(None, None, None))
|
||||
else:
|
||||
indexes = (
|
||||
(indexes[0], 0, indexes[1])
|
||||
if len(self.original_ys_shape) == 1
|
||||
else (indexes[0], indexes[1], 0)
|
||||
)
|
||||
|
||||
indexes = _multi_index(indexes, self.inner_shape)
|
||||
|
||||
if isinstance(indexes[self._lazy_axis], int):
|
||||
other_indexes = (
|
||||
indexes[: self._lazy_axis] + indexes[self._lazy_axis + 1 :]
|
||||
)
|
||||
return self._cached_evaluate(indexes[self._lazy_axis])[
|
||||
other_indexes
|
||||
]
|
||||
lazy_indexes = self._lazy_indexes(indexes)
|
||||
# Using concat and reshape to replace stack operator temporarily, as
|
||||
# it is not a primitive operator.
|
||||
shape = list(self.inner_shape)
|
||||
shape[self._lazy_axis] = len(lazy_indexes)
|
||||
part_jac = paddle.concat(
|
||||
[self._cached_evaluate(i) for i in lazy_indexes],
|
||||
axis=self._lazy_axis,
|
||||
).reshape(shape)
|
||||
result = part_jac[self._shifted_indexes(indexes, len(lazy_indexes))]
|
||||
|
||||
# squeeze redundant 1 in shape
|
||||
if len(result.shape) > len(self.shape):
|
||||
for _ in range(len(result.shape) - len(self.shape)):
|
||||
result = result.squeeze(-1)
|
||||
|
||||
return result
|
||||
|
||||
def _cached_evaluate(self, k):
|
||||
if k is None:
|
||||
return self._cached_evaluate(0).reshape([])
|
||||
v = self._cache.get(k)
|
||||
if v is None:
|
||||
v = self._evaluate(k)
|
||||
self._cache[k] = v
|
||||
return v
|
||||
|
||||
def _evaluate(self, index):
|
||||
"""Evaluate one slice at along lazy axis."""
|
||||
raise NotImplementedError
|
||||
|
||||
def _evaluate_all(self):
|
||||
if len(self.shape) == 0:
|
||||
return self._cached_evaluate(None)
|
||||
else:
|
||||
return self[:]
|
||||
|
||||
|
||||
class _JacobianNoBatch(_Jacobian):
|
||||
"""Compute Jacobian matrix without batch dimension.
|
||||
Suppose the mapping is :math:`f: R^M \to R^N`, the output shape is
|
||||
``(N, M)`` .
|
||||
"""
|
||||
|
||||
def __init__(self, ys, xs):
|
||||
self.is_batched = False
|
||||
super().__init__(ys, xs)
|
||||
# inner_shape is for convenient, it will regard 0-D tensor as 1-D tensor
|
||||
self.inner_shape = [
|
||||
*(self._flatten_ys.shape[0:1]),
|
||||
*(self._flatten_xs.shape[0:1]),
|
||||
]
|
||||
self.shape = [
|
||||
*(self.original_ys_shape[0:1]),
|
||||
*(self.original_xs_shape[0:1]),
|
||||
]
|
||||
|
||||
@property
|
||||
def _lazy_axis(self):
|
||||
return 0
|
||||
|
||||
def _flatten(self, xs):
|
||||
if not isinstance(xs, Sequence):
|
||||
return xs.reshape((-1,))
|
||||
return paddle.concat(tuple(x.reshape((-1,)) for x in xs))
|
||||
|
||||
def _evaluate(self, row_index):
|
||||
return self._flatten(
|
||||
_grad_for_jacobian(
|
||||
self._flatten_ys[row_index],
|
||||
self._xs,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
class _JacobianBatchFirst(_Jacobian):
|
||||
"""Compute Jacobian matrix with batch at first axis.
|
||||
Suppose the mapping is :math:`f: R^{B,M} \to R^{B,N}`, the output shape is
|
||||
``(B, N, M)`` .
|
||||
"""
|
||||
|
||||
def __init__(self, ys, xs):
|
||||
self.is_batched = True
|
||||
super().__init__(ys, xs)
|
||||
# inner_shape is for convenient, it will regard 0-D tensor as 1-D tensor
|
||||
self.inner_shape = [
|
||||
*(self._flatten_xs.shape[0:1]),
|
||||
*(self._flatten_ys.shape[1:2]),
|
||||
*(self._flatten_xs.shape[1:2]),
|
||||
]
|
||||
self.shape = [
|
||||
*(self._flatten_xs.shape[0:1]),
|
||||
*(self.original_ys_shape[1:2]),
|
||||
*(self.original_xs_shape[1:2]),
|
||||
]
|
||||
|
||||
@property
|
||||
def _lazy_axis(self):
|
||||
return 1
|
||||
|
||||
def _flatten(self, xs):
|
||||
if not isinstance(xs, Sequence):
|
||||
return xs.reshape((xs.shape[0], -1))
|
||||
return paddle.concat(
|
||||
tuple(x.reshape((x.shape[0], -1)) for x in as_tensors(xs)), 1
|
||||
)
|
||||
|
||||
def _evaluate(self, row_index):
|
||||
return self._flatten(
|
||||
_grad_for_jacobian(self._flatten_ys[:, row_index], self._xs)
|
||||
)
|
||||
|
||||
|
||||
def _multi_index(indexes, shape):
|
||||
"""A tool for parsing N-dimensional index into a standard format.
|
||||
|
||||
Currently supporting following input format:
|
||||
* ([positive|negative|slice], ...), the right-most elements can be
|
||||
omitted.
|
||||
|
||||
The standard format after converted is slice tuple which contains N elements:
|
||||
* ([positive|slice], ..., [positive|slice])
|
||||
|
||||
Notes:
|
||||
Ellipsis indexes such as ``(..., i), (i, ...)`` is not supported.
|
||||
|
||||
Args:
|
||||
indexes (tuple): The input indexes.
|
||||
shape (tuple): The input shape.
|
||||
|
||||
Returns:
|
||||
tuple: The standard format index as the above description.
|
||||
"""
|
||||
indexes = indexes if isinstance(indexes, Sequence) else (indexes,)
|
||||
if any(isinstance(i, type(Ellipsis)) for i in indexes):
|
||||
raise IndexError('Ellipsis index currently is not supported.')
|
||||
# Fill the right-most elements.
|
||||
indexes = indexes + (slice(0, None, None),) * (len(shape) - len(indexes))
|
||||
# Convert to positive index.
|
||||
positive_indexes = []
|
||||
for i, index in enumerate(indexes):
|
||||
if isinstance(index, slice):
|
||||
index = slice(
|
||||
index.start or 0, index.stop or shape[i], index.step or 1
|
||||
)
|
||||
positive_indexes.append(
|
||||
slice(
|
||||
index.start + shape[i] if index.start < 0 else index.start,
|
||||
index.stop + shape[i] if index.stop < 0 else index.stop,
|
||||
# Negative step means index backward, no need to convert to
|
||||
# positive integer.
|
||||
index.step,
|
||||
)
|
||||
)
|
||||
elif isinstance(index, int):
|
||||
positive_indexes.append(index + shape[i] if index < 0 else index)
|
||||
else:
|
||||
raise TypeError(f'Not supported index type {index}.')
|
||||
return tuple(positive_indexes)
|
||||
|
||||
|
||||
@overload
|
||||
def jacobian(
|
||||
ys: Tensor,
|
||||
xs: Tensor,
|
||||
batch_axis: int | None = ...,
|
||||
) -> Jacobian: ...
|
||||
|
||||
|
||||
@overload
|
||||
def jacobian(
|
||||
ys: Sequence[Tensor],
|
||||
xs: Sequence[Tensor],
|
||||
batch_axis: int | None = ...,
|
||||
) -> tuple[tuple[Jacobian, ...], ...]: ...
|
||||
|
||||
|
||||
@overload
|
||||
def jacobian(
|
||||
ys: Tensor,
|
||||
xs: Sequence[Tensor],
|
||||
batch_axis: int | None = ...,
|
||||
) -> tuple[Jacobian, ...]: ...
|
||||
|
||||
|
||||
@overload
|
||||
def jacobian(
|
||||
ys: Sequence[Tensor],
|
||||
xs: Tensor,
|
||||
batch_axis: int | None = ...,
|
||||
) -> tuple[Jacobian, ...]: ...
|
||||
|
||||
|
||||
def jacobian(
|
||||
ys,
|
||||
xs,
|
||||
batch_axis=None,
|
||||
):
|
||||
r"""
|
||||
Computes the Jacobian of the dependent variable ``ys`` versus the independent
|
||||
variable ``xs``.
|
||||
|
||||
Where ``ys`` represents the output of ``xs`` after a certain operation, ``ys`` and
|
||||
``xs`` can be Tensor or tuple of Tensors, ``batch_axis`` indicates the position of
|
||||
the batch dimension of the parameter data.
|
||||
|
||||
When the input is a tuple Tensors, the returned result is a ``Jacobian`` object with
|
||||
the same number of nesting levels as ``xs``, and each Jacobian has the same shape as
|
||||
The ``xs`` tuples are identical in one-to-one correspondence.
|
||||
|
||||
- When ``batch_axis=None``, only 0-dimensional Tensor or 1-dimensional Tensor is
|
||||
supported, assuming the shape of ``xs`` is ``[N, ]``, the shape of ``ys`` is
|
||||
``[M, ]``, then the output Jacobian matrix shape is ``[M, N]``.
|
||||
|
||||
- When ``batch_axis=0``, only 1-dimensional Tensor or 2-dimensional Tensor is
|
||||
supported, assuming the shape of ``xs`` is ``[B, N]``, The shape of ``ys`` is
|
||||
``[B, M]``, then the output Jacobian matrix shape is ``[B, M, N]``.
|
||||
|
||||
After the ``Jacobian`` object is created, the actual calculation process does not
|
||||
occur, but the lazy evaluation method is used for calculation. It can be
|
||||
multi-dimensional indexed to obtain the entire Jacobian matrix or sub-matrix, and
|
||||
the actual calculation will be performed at this time the value is calculated and
|
||||
the result is returned. At the same time, in the actual evaluation process, the
|
||||
calculated sub-matrix will be cached to avoid duplicate calculations in the
|
||||
subsequent indexing process.
|
||||
|
||||
For example, assuming ``Jacobian`` instance ``J`` has shape ``[B, M, N]``, assuming
|
||||
``M > 4`` , then ``J[:, 1:4:1, :]`` means to get the values from row ``1`` to row
|
||||
``3`` of ``J``. In actual calculation, only the rows ``1`` to ``3`` are evaluated,
|
||||
and the calculation results of ``1`` to ``3`` will be cached at the granularity of
|
||||
the row, and will be used next time. When obtaining one or more rows of results
|
||||
above, the already calculated parts will not be recalculated.
|
||||
|
||||
Args:
|
||||
|
||||
ys (Union[paddle.Tensor, Tuple[paddle.Tensor, ...]]): Output or tuple of outputs derived from xs.
|
||||
xs (Union[paddle.Tensor, Tuple[paddle.Tensor, ...]]): Input or tuple of inputs.
|
||||
batch_axis (Optional[int], optional): Index of batch axis. Defaults to None.
|
||||
|
||||
Returns:
|
||||
|
||||
Union[Tuple[Tuple[Jacobian, ...], ...], Tuple[Jacobian, ...], Jacobian]: Jacobian(s) of ys derived from xs.
|
||||
|
||||
Examples:
|
||||
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> import paddle
|
||||
|
||||
>>> x1 = paddle.randn([3])
|
||||
>>> x2 = paddle.randn([3])
|
||||
>>> x1.stop_gradient = False
|
||||
>>> x2.stop_gradient = False
|
||||
|
||||
>>> y = x1 + x2
|
||||
|
||||
>>> J = paddle.autograd.jacobian(y, (x1, x2))
|
||||
>>> J_y_x1 = J[0][:] # evaluate result of dy/dx1
|
||||
>>> J_y_x2 = J[1][:] # evaluate result of dy/dx2
|
||||
|
||||
>>> print(J_y_x1.shape)
|
||||
paddle.Size([3, 3])
|
||||
>>> print(J_y_x2.shape)
|
||||
paddle.Size([3, 3])
|
||||
"""
|
||||
|
||||
if batch_axis is not None and batch_axis != 0:
|
||||
raise ValueError(
|
||||
f"batch_axis should be None or 0, but got {batch_axis}."
|
||||
)
|
||||
|
||||
# TODO(HydrogenSulfate): support batch_axis > 0
|
||||
is_batched = batch_axis is not None
|
||||
if isinstance(ys, Sequence) and isinstance(xs, Sequence):
|
||||
_jacobian = tuple(
|
||||
tuple(Jacobian(_ys, _xs, is_batched) for _xs in xs) for _ys in ys
|
||||
)
|
||||
elif isinstance(ys, Sequence) and not isinstance(xs, Sequence):
|
||||
_jacobian = tuple(Jacobian(_ys, xs, is_batched) for _ys in ys)
|
||||
elif not isinstance(ys, Sequence) and isinstance(xs, Sequence):
|
||||
_jacobian = tuple(Jacobian(ys, _xs, is_batched) for _xs in xs)
|
||||
else:
|
||||
_jacobian = Jacobian(ys, xs, is_batched)
|
||||
|
||||
return _jacobian
|
||||
|
||||
|
||||
@overload
|
||||
def hessian(
|
||||
ys: Tensor,
|
||||
xs: Tensor,
|
||||
batch_axis: int | None = ...,
|
||||
) -> Hessian: ...
|
||||
|
||||
|
||||
@overload
|
||||
def hessian(
|
||||
ys: Tensor,
|
||||
xs: Sequence[Tensor],
|
||||
batch_axis: int | None = ...,
|
||||
) -> tuple[tuple[Hessian, ...], ...]: ...
|
||||
|
||||
|
||||
def hessian(
|
||||
ys,
|
||||
xs,
|
||||
batch_axis=None,
|
||||
):
|
||||
r"""
|
||||
Computes the Jacobian of the dependent variable ``ys`` versus the independent
|
||||
variable ``xs``.
|
||||
|
||||
Among them, ``ys`` means the output of ``xs`` after a certain operation, ``ys`` can
|
||||
only be a single Tensor, ``xs`` can be a Tensor or a Tensor tuple, and
|
||||
``batch_axis`` means The position of the batch dimension of the parameter data.
|
||||
|
||||
When the input ``xs`` is a Tensor tuple, the returned result is a ``Hessian`` tuple,
|
||||
assuming that the internal shape of the ``xs`` tuple is composed of ``([M1, ], [M2, ])``, the shape of the returned
|
||||
result consists of ``(([M1, M1], [M1, M2]), ([M2, M1], [M2, M2]))``
|
||||
|
||||
- When ``batch_axis=None``, only 0-dimensional Tensor or 1-dimensional Tensor is
|
||||
supported, assuming that the shape of ``xs`` is ``[N, ]``, and the shape of ``ys`` is ``[ ]`` (0-dimensional Tensor), the final output is a single Hessian matrix whose shape is ``[N, N]``.
|
||||
|
||||
- When ``batch_axis=0``, only 1-dimensional Tensor or 2-dimensional Tensor is
|
||||
supported, assuming that the shape of ``xs`` is ``[B, N]``, and the shape of ``ys`` is ``[B, ]``, the final output Jacobian matrix shape is ``[B, N, N]``.
|
||||
|
||||
After the ``Hessian`` object is created, the complete calculation process does not
|
||||
occur, but a partial lazy evaluation method is used for calculation. It can be
|
||||
multi-dimensionally indexed to obtain the entire Hessian matrix or sub-matrix. At
|
||||
this time, the actual Evaluates the computation and returns the result. At the same
|
||||
time, in the actual evaluation process, the calculated sub-matrix will be cached to
|
||||
avoid repeated calculations in the subsequent indexing process.
|
||||
|
||||
Args:
|
||||
|
||||
ys (paddle.Tensor): Output derived from xs which contain one element.
|
||||
xs (Union[paddle.Tensor, Tuple[paddle.Tensor, ...]]): Input or tuple of inputs.
|
||||
batch_axis (Optional[int], optional): Index of batch axis. Defaults to None.
|
||||
|
||||
Returns:
|
||||
|
||||
Union[Tuple[Tuple[Hessian, ...], ...], Tuple[Hessian, ...], Hessian]: Hessian(s) of ys derived from xs.
|
||||
|
||||
Examples:
|
||||
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> import paddle
|
||||
|
||||
>>> x1 = paddle.randn([3])
|
||||
>>> x2 = paddle.randn([4])
|
||||
>>> x1.stop_gradient = False
|
||||
>>> x2.stop_gradient = False
|
||||
|
||||
>>> y = x1.sum() + x2.sum()
|
||||
|
||||
>>> H = paddle.autograd.hessian(y, (x1, x2))
|
||||
>>> H_y_x1_x1 = H[0][0][:] # evaluate result of ddy/dx1x1
|
||||
>>> H_y_x1_x2 = H[0][1][:] # evaluate result of ddy/dx1x2
|
||||
>>> H_y_x2_x1 = H[1][0][:] # evaluate result of ddy/dx2x1
|
||||
>>> H_y_x2_x2 = H[1][1][:] # evaluate result of ddy/dx2x2
|
||||
|
||||
>>> print(H_y_x1_x1.shape)
|
||||
paddle.Size([3, 3])
|
||||
>>> print(H_y_x1_x2.shape)
|
||||
paddle.Size([3, 4])
|
||||
>>> print(H_y_x2_x1.shape)
|
||||
paddle.Size([4, 3])
|
||||
>>> print(H_y_x2_x2.shape)
|
||||
paddle.Size([4, 4])
|
||||
"""
|
||||
|
||||
if batch_axis is None:
|
||||
if ys.numel() > 1:
|
||||
raise ValueError(
|
||||
f"Only support ys.numel()({ys.numel()})==1 when batch_axis is None."
|
||||
)
|
||||
ys = ys.reshape(())
|
||||
elif isinstance(batch_axis, int):
|
||||
if ys[0].numel() > 1:
|
||||
raise ValueError(
|
||||
f"Only support ys[0].numel()({ys.numel()})==1 when batch_axis is int"
|
||||
)
|
||||
# TODO(HydrogenSulfate): support batch_axis > 0
|
||||
if batch_axis != 0:
|
||||
raise ValueError("Only support batch_axis=0 yet.")
|
||||
ys = ys.reshape((-1,))
|
||||
else:
|
||||
raise ValueError(
|
||||
f"batch_axis should be None or int, but got {type(batch_axis)}."
|
||||
)
|
||||
|
||||
_jacobian = jacobian(ys, xs, batch_axis)
|
||||
if not isinstance(xs, Sequence):
|
||||
hessian = jacobian(_jacobian, xs, batch_axis)
|
||||
|
||||
# change classname to Hessian instead of Jacobian.
|
||||
hessian.__class__ = Hessian
|
||||
else:
|
||||
hessian = tuple(jacobian(_j, xs, batch_axis) for _j in _jacobian)
|
||||
|
||||
# change classname to Hessian instead of Jacobian.
|
||||
for i in range(len(hessian)):
|
||||
for j in range(len(hessian[0])):
|
||||
hessian[i][j].__class__ = Hessian
|
||||
|
||||
return hessian
|
||||
|
||||
|
||||
def _replace_none_with_zero_tensor(xs, refs):
|
||||
if xs is None:
|
||||
xs = paddle.zeros_like(refs)
|
||||
xs.stop_gradient = refs.stop_gradient
|
||||
return xs
|
||||
elif isinstance(xs, Sequence):
|
||||
return tuple(
|
||||
_replace_none_with_zero_tensor(x, refs[i]) for i, x in enumerate(xs)
|
||||
)
|
||||
else:
|
||||
return xs
|
||||
|
||||
|
||||
def _grad_for_jacobian(ys, xs, v=None):
|
||||
"""A gradient function that can be used in dynamic graph and static graph.
|
||||
|
||||
The ``grad`` combines ``paddle.grad`` used in dynamic graph and
|
||||
``paddle.static.gradients`` used in static graph, and do following changes:
|
||||
|
||||
* The ``allow_unused`` flag is removed and set defaults to true internally,
|
||||
none in outputs will be replaced by zero tensor.
|
||||
* The ``create_graph`` flag is removed and set defaults to true internally,
|
||||
only makes sense in dynamic graph.
|
||||
* When xs is a single Tensor, ``paddle.grad`` returns a list which only
|
||||
contains one Tensor. It may confuse users, thus in this case we improve
|
||||
to return a single Tensor in _grad_for_jacobian interface.
|
||||
|
||||
Args:
|
||||
ys (Tensor|Sequence[Tensor]): The output tensor or tensor sequence of
|
||||
the graph to compute gradients.
|
||||
xs (Tensor|Sequence[Tensor]): The input tensor or tensor sequence of the graph to
|
||||
compute gradients. The returned values of this API are the
|
||||
gradients of inputs .
|
||||
v (Tensor|Sequence[Tensor]|None,optional): The initial gradient values
|
||||
of outputs . If grad_outputs is None, the initial gradient values of
|
||||
outputs would be Tensors filled with 1; if grad_outputs is not None,
|
||||
it must have the same length as outputs , and in this case, the
|
||||
initial gradient value of the i-th outputs would be: (1) a Tensor
|
||||
filled with 1 when the i-th element of grad_outputs is None;
|
||||
(2) the i-th element of grad_outputs when the i-th element of
|
||||
grad_outputs is a Tensor. Default None.
|
||||
|
||||
Returns:
|
||||
Tensor|tuple[Tensor]: Tensor or a tuple of Tensors, whose length is the
|
||||
same as the Tensor number inside inputs, and the i-th returned
|
||||
Tensor is the sum of gradients of outputs with respect to the i-th
|
||||
inputs.
|
||||
"""
|
||||
if paddle.in_dynamic_mode():
|
||||
# paddle.grad returns a list though the inputs is a single Tensor. The
|
||||
# follow code snippet fixes the problem by return the first element of
|
||||
# xs_grad when the xs is a single Tensor.
|
||||
xs_grad = paddle.grad(ys, xs, v, create_graph=True, allow_unused=True)
|
||||
if (
|
||||
isinstance(xs, paddle.base.framework.Variable)
|
||||
and isinstance(xs_grad, Sequence)
|
||||
and len(xs_grad) > 0
|
||||
):
|
||||
xs_grad = xs_grad[0]
|
||||
else:
|
||||
xs_grad = paddle.static.gradients(ys, xs, v)
|
||||
if (
|
||||
isinstance(xs, framework.Variable)
|
||||
and isinstance(xs_grad, Sequence)
|
||||
and len(xs_grad) > 0
|
||||
):
|
||||
xs_grad = xs_grad[0]
|
||||
return _replace_none_with_zero_tensor(xs_grad, xs)
|
||||
@@ -0,0 +1,152 @@
|
||||
# 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
|
||||
|
||||
import paddle
|
||||
from paddle.base import core, framework
|
||||
from paddle.base.backward import gradients_with_optimizer # noqa: F401
|
||||
from paddle.utils.download import check_and_create_dir
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Sequence
|
||||
|
||||
from paddle import Tensor
|
||||
|
||||
|
||||
__all__ = []
|
||||
|
||||
|
||||
@framework.dygraph_only
|
||||
def backward(
|
||||
tensors: Tensor | Sequence[Tensor],
|
||||
grad_tensors: Tensor | Sequence[Tensor | None] | None = None,
|
||||
retain_graph: bool = False,
|
||||
create_graph: bool = False,
|
||||
*,
|
||||
dump_backward_graph_path: str | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
Compute the backward gradients of given tensors.
|
||||
|
||||
Args:
|
||||
tensors(list of Tensors): the tensors which the gradient to be computed. The tensors can not contain the same tensor.
|
||||
|
||||
grad_tensors(list of Tensors of None, optional): the init gradients of the `tensors`` .If not None, it must have the same length with ``tensors`` ,
|
||||
and if any of the elements is None, then the init gradient is the default value which is filled with 1.0.
|
||||
If None, all the gradients of the ``tensors`` is the default value which is filled with 1.0.
|
||||
Defaults to None.
|
||||
|
||||
retain_graph(bool, optional): If False, the graph used to compute grads will be freed. If you would
|
||||
like to add more ops to the built graph after calling this method( :code:`backward` ), set the parameter
|
||||
:code:`retain_graph` to True, then the grads will be retained. Thus, setting it to False is much more memory-efficient.
|
||||
Defaults to False.
|
||||
dump_backward_graph_path(str, optional): Specifies the directory path for storing the debug file.
|
||||
If this parameter is specified, the backward-related graph (in dot format)
|
||||
and the debugging call stack information will be generated in this directory.
|
||||
Returns:
|
||||
NoneType: None
|
||||
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> import paddle
|
||||
>>> x = paddle.to_tensor([[1, 2], [3, 4]], dtype='float32', stop_gradient=False)
|
||||
>>> y = paddle.to_tensor([[3, 2], [3, 4]], dtype='float32')
|
||||
|
||||
>>> grad_tensor1 = paddle.to_tensor([[1, 2], [2, 3]], dtype='float32')
|
||||
>>> grad_tensor2 = paddle.to_tensor([[1, 1], [1, 1]], dtype='float32')
|
||||
|
||||
>>> z1 = paddle.matmul(x, y)
|
||||
>>> z2 = paddle.matmul(x, y)
|
||||
|
||||
>>> paddle.autograd.backward([z1, z2], [grad_tensor1, grad_tensor2], True)
|
||||
>>> print(x.grad)
|
||||
Tensor(shape=[2, 2], dtype=float32, place=Place(cpu), stop_gradient=False,
|
||||
[[12., 18.],
|
||||
[17., 25.]])
|
||||
|
||||
|
||||
>>> x.clear_grad()
|
||||
|
||||
>>> paddle.autograd.backward([z1, z2], [grad_tensor1, None], True)
|
||||
>>> print(x.grad)
|
||||
Tensor(shape=[2, 2], dtype=float32, place=Place(cpu), stop_gradient=False,
|
||||
[[12., 18.],
|
||||
[17., 25.]])
|
||||
|
||||
>>> x.clear_grad()
|
||||
|
||||
>>> paddle.autograd.backward([z1, z2])
|
||||
>>> print(x.grad)
|
||||
Tensor(shape=[2, 2], dtype=float32, place=Place(cpu), stop_gradient=False,
|
||||
[[10., 14.],
|
||||
[10., 14.]])
|
||||
|
||||
|
||||
"""
|
||||
|
||||
def check_tensors(
|
||||
in_out_list: Sequence[Tensor] | Tensor, name: str
|
||||
) -> Sequence[Tensor]:
|
||||
assert in_out_list is not None, f"{name} should not be None"
|
||||
|
||||
if isinstance(in_out_list, (list, tuple)):
|
||||
assert len(in_out_list) > 0, f"{name} cannot be empty"
|
||||
for each_var in in_out_list:
|
||||
assert isinstance(each_var, paddle.Tensor), (
|
||||
f"Elements of {name} must be paddle.Tensor"
|
||||
)
|
||||
return in_out_list
|
||||
else:
|
||||
assert isinstance(in_out_list, paddle.Tensor), (
|
||||
f"{name} must be Tensor or list of Tensor"
|
||||
)
|
||||
return [in_out_list]
|
||||
|
||||
tensors = check_tensors(tensors, "tensors")
|
||||
|
||||
assert len(tensors) == len(set(tensors)), (
|
||||
"The argument 'tensors' of paddle.autograd.backward contains duplicate paddle.Tensor object."
|
||||
)
|
||||
|
||||
if grad_tensors is not None:
|
||||
if not isinstance(grad_tensors, (list, tuple)):
|
||||
grad_tensors = [grad_tensors]
|
||||
|
||||
for each_tensor in grad_tensors:
|
||||
if each_tensor is not None:
|
||||
assert isinstance(each_tensor, paddle.Tensor), (
|
||||
"The argument 'grad_tensors' of paddle.autograd.backward is invalid, it can be 'None', 'paddle.Tensor' or 'list[None/paddle.Tensor]'."
|
||||
)
|
||||
else:
|
||||
grad_tensors = []
|
||||
|
||||
if len(grad_tensors) > 0:
|
||||
assert len(tensors) == len(grad_tensors), (
|
||||
"The length of grad_tensors must be equal to tensors"
|
||||
)
|
||||
|
||||
assert isinstance(retain_graph, bool), "retain_graph must be True or False"
|
||||
check_and_create_dir(dump_backward_graph_path)
|
||||
core.eager.run_backward(
|
||||
tensors,
|
||||
grad_tensors,
|
||||
retain_graph,
|
||||
create_graph,
|
||||
dump_backward_graph_path,
|
||||
)
|
||||
@@ -0,0 +1,810 @@
|
||||
# Copyright (c) 2023 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,tes
|
||||
# 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
|
||||
|
||||
import collections
|
||||
import logging
|
||||
import warnings
|
||||
from collections.abc import Sequence
|
||||
from functools import lru_cache
|
||||
from typing import Any
|
||||
|
||||
from paddle import pir
|
||||
from paddle.base import core
|
||||
from paddle.base.libpaddle.pir import (
|
||||
get_used_external_value,
|
||||
)
|
||||
from paddle.base.wrapped_decorator import signature_safe_contextmanager
|
||||
|
||||
# TODO(CZ): to be removed when we support dynamic shape by default.
|
||||
ALLOW_DYNAMIC_SHAPE_VJP_OPS = [
|
||||
"pd_op.abs",
|
||||
"pd_op.add",
|
||||
"pd_op.amax",
|
||||
"pd_op.amin",
|
||||
"pd_op.angle",
|
||||
"pd_op.argsort",
|
||||
"pd_op.assign",
|
||||
"pd_op.batch_norm_",
|
||||
"pd_op.cast",
|
||||
"pd_op.ceil",
|
||||
"pd_op.concat",
|
||||
"pd_op.cos",
|
||||
"pd_op.cumprod",
|
||||
"pd_op.cumsum",
|
||||
"pd_op.divide",
|
||||
"pd_op.dot",
|
||||
"pd_op.dropout",
|
||||
"pd_op.elementwise_pow",
|
||||
"pd_op.erf",
|
||||
"pd_op.exp",
|
||||
"pd_op.expand",
|
||||
"pd_op.floor",
|
||||
"pd_op.fmax",
|
||||
"pd_op.fmin",
|
||||
"pd_op.gather",
|
||||
"pd_op.gather_nd",
|
||||
"pd_op.gelu",
|
||||
"pd_op.group_norm",
|
||||
"pd_op.hardsigmoid",
|
||||
"pd_op.hardswish",
|
||||
"pd_op.kron",
|
||||
"pd_op.kthvalue",
|
||||
"pd_op.layer_norm",
|
||||
"pd_op.leaky_relu",
|
||||
"pd_op.log",
|
||||
"pd_op.logcumsumexp",
|
||||
"pd_op.logsumexp",
|
||||
"pd_op.linear_v2",
|
||||
"pd_op.matmul",
|
||||
"pd_op.max",
|
||||
"pd_op.maximum",
|
||||
"pd_op.mean",
|
||||
"pd_op.minimum",
|
||||
"pd_op.multiply",
|
||||
"pd_op.pad",
|
||||
"pd_op.pow",
|
||||
"pd_op.prod",
|
||||
"pd_op.reduce_as",
|
||||
"pd_op.relu",
|
||||
"pd_op.relu6",
|
||||
"pd_op.reshape",
|
||||
"pd_op.roll",
|
||||
"pd_op.rsqrt",
|
||||
"pd_op.scale",
|
||||
"pd_op.scatter",
|
||||
"pd_op.scatter_nd_add",
|
||||
"pd_op.sigmoid",
|
||||
"pd_op.silu",
|
||||
"pd_op.sin",
|
||||
"pd_op.softmax",
|
||||
"pd_op.softsign",
|
||||
"pd_op.split",
|
||||
"pd_op.sqrt",
|
||||
"pd_op.square",
|
||||
"pd_op.squeeze",
|
||||
"pd_op.stack",
|
||||
"pd_op.subtract",
|
||||
"pd_op.sum",
|
||||
"pd_op.swiglu",
|
||||
"pd_op.swish",
|
||||
"pd_op.take_along_axis",
|
||||
"pd_op.tanh",
|
||||
"pd_op.tile",
|
||||
"pd_op.topk",
|
||||
"pd_op.transpose",
|
||||
"pd_op.trunc",
|
||||
"pd_op.unsqueeze",
|
||||
"pd_op.where",
|
||||
"pd_op.p_norm",
|
||||
"pd_op.index_put",
|
||||
"pd_op.index_add",
|
||||
"pd_op.elu",
|
||||
"pd_op.masked_fill",
|
||||
"pd_op.masked_select",
|
||||
"pd_op.var",
|
||||
]
|
||||
|
||||
|
||||
class ValueWrapper:
|
||||
def __init__(self, value) -> None:
|
||||
if isinstance(value, ValueWrapper):
|
||||
assert isinstance(value._value, (type(None), pir.Value))
|
||||
else:
|
||||
if not isinstance(value, (type(None), pir.Value)):
|
||||
raise TypeError(
|
||||
"Value Wrapper is only support None and pir.Value"
|
||||
)
|
||||
self._value = value._value if isinstance(value, ValueWrapper) else value
|
||||
|
||||
def __hash__(self) -> int:
|
||||
if isinstance(self._value, pir.Value):
|
||||
return self._value.hash()
|
||||
else:
|
||||
return hash(self._value)
|
||||
|
||||
def __eq__(self, other) -> bool:
|
||||
if not isinstance(other, ValueWrapper):
|
||||
warnings.warn(
|
||||
f'In ValueWrapper.__eq__ expected type of `other` is ValueWrapper but received {other.__class__}.'
|
||||
)
|
||||
return False
|
||||
|
||||
if self._value is None or other._value is None:
|
||||
return self._value is None and other._value is None
|
||||
return self._value.is_same(other._value)
|
||||
|
||||
|
||||
class ValueDict:
|
||||
def __init__(
|
||||
self,
|
||||
iter=None,
|
||||
*,
|
||||
default_factory=None,
|
||||
):
|
||||
self._items: dict[ValueWrapper] = {}
|
||||
self._default_factory = default_factory
|
||||
if iter is not None:
|
||||
for key, val in iter.items():
|
||||
self[key] = val
|
||||
|
||||
def copy(self):
|
||||
ret = ValueDict()
|
||||
ret._items = self._items.copy()
|
||||
ret._default_factory = self._default_factory
|
||||
return ret
|
||||
|
||||
def update(self, other_dict):
|
||||
for key, val in other_dict.items():
|
||||
self[key] = val
|
||||
|
||||
def keys(self):
|
||||
for key in self._items.keys():
|
||||
yield key._value
|
||||
|
||||
def values(self):
|
||||
return self._items.values()
|
||||
|
||||
def items(self):
|
||||
for key, val in self._items.items():
|
||||
yield key._value, val
|
||||
|
||||
def get(self, key, default=None):
|
||||
if not self.__contains__(key):
|
||||
return default
|
||||
return self._items[ValueWrapper(key)]
|
||||
|
||||
def pop(self, key):
|
||||
if not self.__contains__(key):
|
||||
raise KeyError(f'{key} is not in ValueDict')
|
||||
return self._items.pop(ValueWrapper(key))
|
||||
|
||||
def setdefault(self, key, default=None):
|
||||
if not self.__contains__(key):
|
||||
self[key] = default
|
||||
return self[key]
|
||||
|
||||
def __setitem__(self, key, val: Any):
|
||||
self._items[ValueWrapper(key)] = val
|
||||
|
||||
def __getitem__(self, key):
|
||||
if not self.__contains__(key):
|
||||
if self._default_factory is not None:
|
||||
self[key] = self._default_factory()
|
||||
else:
|
||||
raise KeyError(f'{key} is not in ValueDict')
|
||||
return self._items[ValueWrapper(key)]
|
||||
|
||||
def __bool__(self):
|
||||
return bool(self._items)
|
||||
|
||||
def __len__(self):
|
||||
return len(self._items)
|
||||
|
||||
def __iter__(self):
|
||||
return self.keys()
|
||||
|
||||
def __contains__(self, key):
|
||||
return ValueWrapper(key) in self._items
|
||||
|
||||
def __repr__(self) -> str:
|
||||
items_str = ", ".join(f"{key}: {val}" for key, val in self.items())
|
||||
return f'ValueDict({items_str})'
|
||||
|
||||
|
||||
class ValueSet:
|
||||
def __init__(
|
||||
self, iter: Sequence[ValueWrapper] | set[ValueWrapper] | None = None
|
||||
):
|
||||
self._set: set[ValueWrapper] = set()
|
||||
if iter is not None:
|
||||
for val in iter:
|
||||
self.add(val)
|
||||
|
||||
def copy(self):
|
||||
ret = ValueSet()
|
||||
ret._set = self._set.copy()
|
||||
return ret
|
||||
|
||||
def add(self, val):
|
||||
if not self.__contains__(val):
|
||||
self._set.add(ValueWrapper(val))
|
||||
|
||||
def update(self, other: set):
|
||||
for val in other:
|
||||
self.add(val)
|
||||
|
||||
def pop(self):
|
||||
return self._set.pop()._value
|
||||
|
||||
def remove(self, val):
|
||||
self._set.remove(ValueWrapper(val))
|
||||
|
||||
def discard(self, val):
|
||||
self._set.discard(ValueWrapper(val))
|
||||
|
||||
def __and__(self, other: ValueSet):
|
||||
return ValueSet(self._set & other._set)
|
||||
|
||||
def __sub__(self, other: ValueSet):
|
||||
return ValueSet(self._set - other._set)
|
||||
|
||||
def __or__(self, other: ValueSet):
|
||||
return ValueSet(self._set | other._set)
|
||||
|
||||
def __bool__(self):
|
||||
return bool(self._set)
|
||||
|
||||
def __len__(self):
|
||||
return len(self._set)
|
||||
|
||||
def __iter__(self):
|
||||
for val in self._set:
|
||||
yield val._value
|
||||
|
||||
def __contains__(self, val):
|
||||
return ValueWrapper(val) in self._set
|
||||
|
||||
def __repr__(self) -> str:
|
||||
items_str = ", ".join(repr(item) for item in self)
|
||||
return f'ValueSet({items_str})'
|
||||
|
||||
|
||||
class State:
|
||||
"""
|
||||
record relationship of forward op/value and backward op/value
|
||||
one state must be binding with a block, if block has parent block,
|
||||
state will include parent block info.
|
||||
|
||||
"""
|
||||
|
||||
def __init__(self, block):
|
||||
self.block = block
|
||||
# value -> list(list(value))
|
||||
self.value_to_valuegrad = ValueDict(default_factory=list)
|
||||
self.value_to_sumvaluegrad = ValueDict(default_factory=list)
|
||||
# operation -> list(operation)
|
||||
self.op_to_opgrad = collections.defaultdict(list)
|
||||
|
||||
# value -> list(value)
|
||||
self.valuegrad_to_value = ValueDict(default_factory=list)
|
||||
self.sumvaluegrad_to_value = ValueDict(default_factory=list)
|
||||
# operation -> list(operation)
|
||||
self.opgrad_to_op = collections.defaultdict(list)
|
||||
# only for controlflow
|
||||
# inside_value is sub block value, which will yield to parent block,
|
||||
# parent block value is outside_value
|
||||
self.inside_value_to_outside_value_map = ValueDict()
|
||||
|
||||
def turn_map(self) -> None:
|
||||
self.valuegrad_to_value = ValueDict(default_factory=list)
|
||||
self.sumvaluegrad_to_value = ValueDict(default_factory=list)
|
||||
self.opgrad_to_op = collections.defaultdict(list)
|
||||
|
||||
for k, v in self.value_to_valuegrad.items():
|
||||
if v != []:
|
||||
for value in v[0]:
|
||||
self.valuegrad_to_value[value] = [k]
|
||||
for k, v in self.value_to_sumvaluegrad.items():
|
||||
if v != []:
|
||||
for value in v[0]:
|
||||
self.sumvaluegrad_to_value[value] = [k]
|
||||
for k, v in self.op_to_opgrad.items():
|
||||
if v != []:
|
||||
self.opgrad_to_op[v[0]] = [k]
|
||||
|
||||
def copy(self, new_block):
|
||||
state = State(new_block)
|
||||
state.value_to_valuegrad = self.value_to_valuegrad.copy()
|
||||
state.value_to_sumvaluegrad = self.value_to_sumvaluegrad.copy()
|
||||
|
||||
# operation -> list(operation)
|
||||
state.op_to_opgrad = self.op_to_opgrad.copy()
|
||||
|
||||
# value -> list(value)
|
||||
state.valuegrad_to_value = self.valuegrad_to_value.copy()
|
||||
state.sumvaluegrad_to_value = self.sumvaluegrad_to_value.copy()
|
||||
# operation -> list(operation)
|
||||
state.opgrad_to_op = self.opgrad_to_op.copy()
|
||||
|
||||
# only for controlflow
|
||||
state.inside_value_to_outside_value_map = (
|
||||
self.inside_value_to_outside_value_map.copy()
|
||||
)
|
||||
|
||||
return state
|
||||
|
||||
|
||||
def _check_vjp_dynamic_shape(op, inputs):
|
||||
for items in inputs:
|
||||
for item in items:
|
||||
if (
|
||||
item.is_dense_tensor_type()
|
||||
and item.initialized()
|
||||
and -1 in item.shape
|
||||
):
|
||||
return True
|
||||
|
||||
|
||||
# Prim currently does not support dynamic shape, when dynamic shape exits in shape of op inputs, prim will be skipped its vjp op.
|
||||
@signature_safe_contextmanager
|
||||
def dynamic_shape_prim_vjp_guard(op, inputs):
|
||||
origin_prim = core._is_bwd_prim_enabled()
|
||||
if op.name() == "cf.tuple_push":
|
||||
skip_prim = True
|
||||
else:
|
||||
skip_prim = (
|
||||
origin_prim
|
||||
and core._enable_prim_skip_dynamic_shape()
|
||||
and _check_vjp_dynamic_shape(op, inputs)
|
||||
and op.name() not in ALLOW_DYNAMIC_SHAPE_VJP_OPS
|
||||
)
|
||||
|
||||
try:
|
||||
if origin_prim and skip_prim:
|
||||
core._set_prim_backward_enabled(False)
|
||||
yield
|
||||
finally:
|
||||
if origin_prim:
|
||||
core._set_prim_backward_enabled(True)
|
||||
|
||||
|
||||
def check_type(input, input_name, expected_type, op_name, extra_message=''):
|
||||
if not isinstance(input, expected_type):
|
||||
raise TypeError(
|
||||
f"The type of '{input_name}' in {op_name} must be {expected_type}, but received {type(input)}. {extra_message}"
|
||||
)
|
||||
|
||||
|
||||
def _as_list(x):
|
||||
if x is None:
|
||||
return []
|
||||
return list(x) if isinstance(x, Sequence) else [x]
|
||||
|
||||
|
||||
def some_in_set(value_list, value_set):
|
||||
return any(v in value_set for v in value_list)
|
||||
|
||||
|
||||
def is_control_flow(op):
|
||||
return op.name() == "pd_op.if" or op.name() == "pd_op.while"
|
||||
|
||||
|
||||
def is_builtin_op(op):
|
||||
dialect_name, opname = op.name().split(".")
|
||||
return dialect_name == "builtin"
|
||||
|
||||
|
||||
def update_no_grad_set_by_stopgradient(block, no_grad_set):
|
||||
for op in block.ops:
|
||||
if is_control_flow(op):
|
||||
for sub_block in op.blocks():
|
||||
update_no_grad_set_by_stopgradient(sub_block, no_grad_set)
|
||||
for value in op.results():
|
||||
if value.stop_gradient and value not in no_grad_set:
|
||||
no_grad_set.add(value)
|
||||
|
||||
|
||||
def get_real_op_inputs(op):
|
||||
if op.name() == "pd_op.if":
|
||||
return get_used_external_value(op)
|
||||
elif op.name() == "pd_op.while":
|
||||
return op.operands_source() + get_used_external_value(
|
||||
op.as_while_op().body()
|
||||
)
|
||||
elif op.name() == "pd_op.pylayer":
|
||||
return get_used_external_value(op)
|
||||
else:
|
||||
return op.operands_source()
|
||||
|
||||
|
||||
def get_real_op_outputs(op):
|
||||
outputs = op.results()
|
||||
if op.name() == "pd_op.array_write_":
|
||||
for x in op.operands():
|
||||
outputs.append(x.source())
|
||||
if op.name() == "pd_op.while":
|
||||
for internal_op in op.as_while_op().body().ops:
|
||||
if internal_op.name() == "pd_op.array_write_":
|
||||
for x in internal_op.operands():
|
||||
outputs.append(x.source())
|
||||
return outputs
|
||||
|
||||
|
||||
def inverse_sort_op(old_ops):
|
||||
'''
|
||||
if topo graph is op1 -> op2 -> op3
|
||||
return [op3, op2, op1]
|
||||
|
||||
'''
|
||||
|
||||
# init pending_count[op] which describes number of
|
||||
# pending edges for its grad_op
|
||||
|
||||
pending_count = collections.defaultdict(int)
|
||||
ops = []
|
||||
[ops.append(x) for x in old_ops if x not in ops]
|
||||
ops_set = set(ops)
|
||||
sorted_list = []
|
||||
for op in ops:
|
||||
for x in get_real_op_inputs(op):
|
||||
if not pir.is_fake_value(x) and x.get_defining_op() in ops_set:
|
||||
pending_count[x.get_defining_op()] += 1
|
||||
|
||||
queue = collections.deque()
|
||||
|
||||
for op in ops:
|
||||
if pending_count[op] == 0:
|
||||
queue.append(op)
|
||||
|
||||
while queue:
|
||||
op = queue.popleft()
|
||||
sorted_list.append(op)
|
||||
for x in get_real_op_inputs(op):
|
||||
x_op = x.get_defining_op()
|
||||
pending_count[x_op] -= 1
|
||||
if pending_count[x_op] == 0:
|
||||
queue.append(x_op)
|
||||
|
||||
if len(sorted_list) != len(ops):
|
||||
raise ValueError(
|
||||
"inverse_sort_op wrong, sorted_list size is not equal to origin_list size"
|
||||
)
|
||||
change_list = []
|
||||
# true %0 = op1, 1% = increment(0%), 3% = op2(0%), tuple_push(%0, 1%, 3%),
|
||||
# no one use 1% so increment be the first op, actually op2 use 1% ,
|
||||
# sorted_list = [increment, op2, op1] should be [op2, increment, op1],
|
||||
# tuple_push(0%) must be forward last op, backward first op, so skip it.
|
||||
for op in reversed(sorted_list):
|
||||
if op.name() == 'pd_op.increment_':
|
||||
idx_1 = sorted_list.index(op)
|
||||
idx_2 = sorted_list.index(op)
|
||||
|
||||
for op_in in reversed(sorted_list[: sorted_list.index(op)]):
|
||||
if (
|
||||
some_in_set(
|
||||
op.operands_source(),
|
||||
ValueSet(get_real_op_inputs(op_in)),
|
||||
)
|
||||
and op_in.name() != "cf.tuple_push"
|
||||
):
|
||||
idx_2 = sorted_list.index(op_in)
|
||||
if idx_1 != idx_2:
|
||||
change_list.append((idx_1, idx_2))
|
||||
for idx_1, idx_2 in change_list:
|
||||
sorted_list[idx_1], sorted_list[idx_2] = (
|
||||
sorted_list[idx_2],
|
||||
sorted_list[idx_1],
|
||||
)
|
||||
|
||||
return sorted_list
|
||||
|
||||
|
||||
def is_inplace_net(op_list):
|
||||
'''
|
||||
when program has inplace op , it's difficult to find the actual pending_count.
|
||||
'''
|
||||
for op in op_list:
|
||||
if op.name() in ["pd_op.array_write_", "pd_op.assign_out_"]:
|
||||
return True
|
||||
if is_control_flow(op):
|
||||
for block in op.blocks():
|
||||
if is_inplace_net(block.ops):
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
|
||||
def remove_op(block, op, state):
|
||||
'''
|
||||
remove op from block
|
||||
'''
|
||||
if state.opgrad_to_op[op] != []:
|
||||
fwd_op = state.opgrad_to_op[op][0]
|
||||
state.op_to_opgrad[fwd_op].remove(op)
|
||||
|
||||
for valuegrad in op.results():
|
||||
if state.valuegrad_to_value[valuegrad] != []:
|
||||
value = state.valuegrad_to_value[valuegrad][0]
|
||||
state.value_to_valuegrad[value] = []
|
||||
|
||||
if value in state.sumvaluegrad_to_value:
|
||||
raise ValueError(
|
||||
f'input_grad in [%s] is value which need to sum {op.name()}'
|
||||
)
|
||||
# NOTE(SigureMo): Ensure access to the op's results before removing it.
|
||||
# Otherwise, the op will be deconstructed and access the num_results
|
||||
# will be undefined behavior, it always cause hanging on the macOS.
|
||||
block.remove_op(op)
|
||||
|
||||
|
||||
def while_prune_check(while_tuple_ops):
|
||||
if len(while_tuple_ops) != 0:
|
||||
for opresult in while_tuple_ops[0].results():
|
||||
if not opresult.use_empty():
|
||||
return False
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def remove_useless_full_like_ops(block, ops, state):
|
||||
'''
|
||||
remove ops which are not in use recursively,
|
||||
|
||||
'''
|
||||
remove_ops = []
|
||||
inverse_ops = inverse_sort_op(list(ops))
|
||||
# from output to input
|
||||
for op in inverse_ops:
|
||||
if op.name() == "pd_op.full_like":
|
||||
if op.result(0).use_empty():
|
||||
full_op = op.operand_source(1).get_defining_op()
|
||||
remove_ops.append(op)
|
||||
remove_ops.append(full_op)
|
||||
elif is_control_flow(op):
|
||||
for sub_block in op.blocks():
|
||||
remove_useless_full_like_ops(sub_block, sub_block.ops, state)
|
||||
|
||||
for op in remove_ops:
|
||||
remove_op(block, op, state)
|
||||
|
||||
|
||||
def all_stop_gradient_true(block):
|
||||
for op in block.ops:
|
||||
for value in op.results():
|
||||
if value.stop_gradient is False:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def all_input_stop_gradient_true(list_of_list):
|
||||
for list_ in list_of_list:
|
||||
for stop_gradient in list_:
|
||||
if stop_gradient is False:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def all_output_grad_none(list_of_list):
|
||||
for list_ in list_of_list:
|
||||
for value in list_:
|
||||
if value is not None:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def op_has_vjp(op):
|
||||
# NOTE(MarioLulab): In PIR mode, even though the `PyLayer` op does
|
||||
# not have a vjp interface, we still need to generate the backward
|
||||
# block based on its registered backward function. To achieve this,
|
||||
# we add more handling logic for `PyLayer` Op in the `call_vjp` function
|
||||
return core.has_vjp(op) or op.name() == "pd_op.pylayer"
|
||||
|
||||
|
||||
def parent_total_ops(block):
|
||||
'''
|
||||
when block is sub_block, forward op should include its parent block ops
|
||||
(sub block nest should Add on demand to avoid block copy)
|
||||
'''
|
||||
total_ops = []
|
||||
if block.parent_block is not None:
|
||||
if block.parent_block.parent_block:
|
||||
total_ops += block.parent_block.parent_block.ops
|
||||
total_ops += block.parent_block.ops
|
||||
total_ops += block.ops
|
||||
|
||||
return total_ops
|
||||
|
||||
|
||||
# only for control_flow to find corresponding value or value_list
|
||||
def return_map_value(value, map):
|
||||
output = value
|
||||
while output in map:
|
||||
output = map[output]
|
||||
return output
|
||||
|
||||
|
||||
def return_map_value_list(value, map):
|
||||
output = []
|
||||
for i in range(len(value)):
|
||||
if value[i] in map:
|
||||
output.append(return_map_value(value[i], map))
|
||||
else:
|
||||
output.append(value[i])
|
||||
return output
|
||||
|
||||
|
||||
def argument_to_value(while_op):
|
||||
'''
|
||||
return while op's relationship of (block_argument to input value) and (input value to block_argument).
|
||||
'''
|
||||
if while_op.name() != "pd_op.while":
|
||||
return ValueDict(), ValueDict()
|
||||
|
||||
assert len(while_op.as_while_op().block_arguments()) + 1 == len(
|
||||
while_op.operands_source()
|
||||
), (
|
||||
"while op's block_arguments size + 1 should same to while op's operands_source size"
|
||||
)
|
||||
arg_to_value_map = ValueDict()
|
||||
value_to_arg_map = ValueDict()
|
||||
for arg, value in zip(
|
||||
while_op.as_while_op().block_arguments(),
|
||||
while_op.operands_source()[1:],
|
||||
):
|
||||
arg_to_value_map[arg] = value
|
||||
value_to_arg_map[value] = arg
|
||||
return arg_to_value_map, value_to_arg_map
|
||||
|
||||
|
||||
def get_grad_semantic_info(op):
|
||||
'''
|
||||
return whether op's inputs has grad, usually handled from yaml.
|
||||
some op has uncertain inputs need special handling.
|
||||
'''
|
||||
if op.name() in [
|
||||
"builtin.combine",
|
||||
"pd_op.if",
|
||||
"pd_op.while",
|
||||
"pd_op.pylayer",
|
||||
"cf.tuple_push",
|
||||
"dist_op.moe_global_mesh_tensor",
|
||||
"dist_op.moe_sub_mesh_tensors",
|
||||
"dist_op.dist_reshape",
|
||||
]:
|
||||
grad_semantic_info = [True for _ in range(len(get_real_op_inputs(op)))]
|
||||
if op.name() == "pd_op.if":
|
||||
grad_semantic_info[0] = False
|
||||
else:
|
||||
grad_semantic_info = op.get_input_grad_semantics()
|
||||
return grad_semantic_info
|
||||
|
||||
|
||||
def get_split_op(value):
|
||||
for op in value.all_used_ops():
|
||||
if op.name() == "builtin.split":
|
||||
return op
|
||||
return None
|
||||
|
||||
|
||||
@lru_cache
|
||||
def warning_once(message: str):
|
||||
logging.warning(message)
|
||||
|
||||
|
||||
def update_if_output_stopgradient(if_op, true_yield_op, false_yield_op):
|
||||
"""
|
||||
Update if_op's stop_gradient based on true_yield_op and false_yield_op.
|
||||
|
||||
Args:
|
||||
true_yield_op: true block of if_op's last op.
|
||||
false_yield_op: false block of if_op's last op.
|
||||
if_op: update it's op_results()'s stop_gradient.
|
||||
"""
|
||||
if (
|
||||
true_yield_op.name() != 'cf.yield'
|
||||
or false_yield_op.name() != 'cf.yield'
|
||||
):
|
||||
raise ValueError("param is not yield op")
|
||||
|
||||
# Check if operands_source sizes match
|
||||
if len(true_yield_op.operands_source()) != len(
|
||||
false_yield_op.operands_source()
|
||||
):
|
||||
raise ValueError("Mismatched yield operands_source sizes")
|
||||
|
||||
# Check if op_results size matches operands_source
|
||||
if len(if_op.results()) != len(true_yield_op.operands_source()):
|
||||
raise ValueError(
|
||||
"Mismatched if op_results size with yield operands_source"
|
||||
)
|
||||
|
||||
# Update if_op's stop_gradient
|
||||
for i in range(len(true_yield_op.operands_source())):
|
||||
stop_grad1 = true_yield_op.operand_source(i).stop_gradient
|
||||
stop_grad2 = false_yield_op.operand_source(i).stop_gradient
|
||||
|
||||
# Set to False if either stop_gradient is False
|
||||
if not stop_grad1 or not stop_grad2:
|
||||
if_op.result(i).stop_gradient = False
|
||||
|
||||
|
||||
def update_while_output_stopgradient(while_op, yield_op):
|
||||
"""
|
||||
Update while_op's stop_gradient based on yield_op.
|
||||
|
||||
Args:
|
||||
yield_op: The yield operation associated with the while loop.
|
||||
while_op: The while operation whose op_results()'s stop_gradient needs to be updated.
|
||||
"""
|
||||
# Check if yield_op is indeed a yield operation
|
||||
if yield_op.name() != 'cf.yield':
|
||||
raise ValueError("yield_op is not a yield operation")
|
||||
|
||||
# Check if operands_source size of yield_op matches op_results size of while_op
|
||||
if len(while_op.results()) + 1 != len(yield_op.operands_source()):
|
||||
raise ValueError(
|
||||
f"Mismatched while op_results size %d with yield operands_source %d. {len(while_op.results()) + 1, len(yield_op.operands_source())}"
|
||||
)
|
||||
|
||||
# Update while_op's stop_gradient
|
||||
for i in range(1, len(yield_op.operands_source())):
|
||||
stop_grad = yield_op.operand_source(i).stop_gradient
|
||||
|
||||
# Set to False if stop_gradient is False
|
||||
if not stop_grad:
|
||||
while_op.result(i - 1).stop_gradient = False
|
||||
|
||||
|
||||
def find_index_of_yield(value, yield_op):
|
||||
for i, v in enumerate(yield_op.operands_source()):
|
||||
if v.is_same(value):
|
||||
return i
|
||||
return -1
|
||||
|
||||
|
||||
def update_tuple_pop_origin_inputs(tuple_pop_outputs):
|
||||
if tuple_pop_outputs == []:
|
||||
return tuple_pop_outputs
|
||||
op = tuple_pop_outputs[0][0].get_defining_op()
|
||||
assert op.name() == "cf.tuple_pop"
|
||||
stack_op = op.operand_source(0).get_defining_op()
|
||||
tuple_push_inputs = stack_op.result(1).first_use().owner().operands_source()
|
||||
tuple_push_inputs_with_if = []
|
||||
for input in tuple_push_inputs:
|
||||
if input.first_use().owner().name() == "cf.yield":
|
||||
yield_op = input.first_use().owner()
|
||||
index = find_index_of_yield(input, yield_op)
|
||||
assert index != -1
|
||||
tuple_push_inputs_with_if.append(
|
||||
yield_op.get_parent_block().parent_op.result(index)
|
||||
)
|
||||
else:
|
||||
tuple_push_inputs_with_if.append(input)
|
||||
|
||||
# pass inlets
|
||||
return tuple_push_inputs_with_if[1:]
|
||||
|
||||
|
||||
def value_in_block(value, block):
|
||||
value_block = value.get_defining_op().get_parent_block()
|
||||
while block.parent_op.name() != "builtin.module":
|
||||
if block == value_block:
|
||||
return True
|
||||
block = block.parent_block
|
||||
# now block is module op's block
|
||||
if block == value_block:
|
||||
return True
|
||||
|
||||
return False
|
||||
@@ -0,0 +1,18 @@
|
||||
# Copyright (c) 2026 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 .py_layer import (
|
||||
PyLayerContext as FunctionCtx, # noqa: F401
|
||||
once_differentiable, # noqa: F401
|
||||
)
|
||||
@@ -0,0 +1,15 @@
|
||||
# Copyright (c) 2025 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 paddle.base.dygraph.base import set_grad_enabled # noqa: F401
|
||||
File diff suppressed because it is too large
Load Diff
@@ -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
|
||||
@@ -0,0 +1,127 @@
|
||||
# Copyright (c) 2022 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
|
||||
|
||||
from paddle.base import core
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Callable
|
||||
|
||||
from paddle import Tensor
|
||||
__all__ = []
|
||||
|
||||
|
||||
class saved_tensors_hooks:
|
||||
"""
|
||||
Dynamic graph, registers a pair of pack / unpack hooks for saved tensors.
|
||||
|
||||
Parameters:
|
||||
pack_hook (function): The pack hook will be called every time the forward
|
||||
operation inputs/outputs tensors need be saved for backward. Then you
|
||||
can save it to CPU or Disk. The input of `pack_hook` is a tensor need
|
||||
be saved. The output of `pack_hook` is then stored information instead
|
||||
of the original tensor. `pack_hook` will also be called while any
|
||||
tensor need be saved by `PyLayerContext.save_for_backward`. If a tensor
|
||||
saved for backward is no need buffer, `pack_hook` will not be called.
|
||||
Only the tensor saved for backward is DenseTensor, `pack_hook` will be
|
||||
called.
|
||||
unpack_hook (function): The unpack hook will be called every time the
|
||||
backward need use the saved inputs/outputs tensors. Then you can reload
|
||||
the tensor and return it to paddle framework. The input of `unpack_hook`
|
||||
is the information returned by `pack_hook`. The output of `unpack_hook`
|
||||
is a tensor reloaded by the information, and the tensor must has the same
|
||||
content as the original tensor passed as input to the corresponding
|
||||
`pack_hook`.
|
||||
|
||||
Returns:
|
||||
None
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
:name: code-example1
|
||||
|
||||
>>> # Example1
|
||||
>>> import paddle
|
||||
|
||||
>>> def pack_hook(x):
|
||||
... print("Packing", x)
|
||||
... return x.numpy()
|
||||
|
||||
>>> def unpack_hook(x):
|
||||
... print("UnPacking", x)
|
||||
... return paddle.to_tensor(x)
|
||||
|
||||
>>> a = paddle.ones([3, 3])
|
||||
>>> b = paddle.ones([3, 3]) * 2
|
||||
>>> a.stop_gradient = False
|
||||
>>> b.stop_gradient = False
|
||||
>>> with paddle.autograd.saved_tensors_hooks(pack_hook, unpack_hook):
|
||||
... y = paddle.multiply(a, b)
|
||||
>>> y.sum().backward()
|
||||
|
||||
.. code-block:: pycon
|
||||
:name: code-example2
|
||||
|
||||
>>> # Example2
|
||||
>>> import paddle
|
||||
>>> from paddle.autograd import PyLayer
|
||||
|
||||
>>> class cus_multiply(PyLayer):
|
||||
... @staticmethod
|
||||
... def forward(ctx, a, b):
|
||||
... y = paddle.multiply(a, b)
|
||||
... ctx.save_for_backward(a, b)
|
||||
... return y
|
||||
...
|
||||
... @staticmethod
|
||||
... def backward(ctx, dy):
|
||||
... a, b = ctx.saved_tensor()
|
||||
... grad_a = dy * a
|
||||
... grad_b = dy * b
|
||||
... return grad_a, grad_b
|
||||
|
||||
>>> def pack_hook(x):
|
||||
... print("Packing", x)
|
||||
... return x.numpy()
|
||||
|
||||
>>> def unpack_hook(x):
|
||||
... print("UnPacking", x)
|
||||
... return paddle.to_tensor(x)
|
||||
|
||||
>>> a = paddle.ones([3, 3])
|
||||
>>> b = paddle.ones([3, 3]) * 2
|
||||
>>> a.stop_gradient = False
|
||||
>>> b.stop_gradient = False
|
||||
>>> with paddle.autograd.saved_tensors_hooks(pack_hook, unpack_hook):
|
||||
... y = cus_multiply.apply(a, b)
|
||||
>>> y.sum().backward()
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
pack_hook: Callable[[Tensor], Any | None],
|
||||
unpack_hook: Callable[[Any], Tensor | None],
|
||||
) -> None:
|
||||
self.pack_hook = pack_hook
|
||||
self.unpack_hook = unpack_hook
|
||||
|
||||
def __enter__(self) -> None:
|
||||
core.eager.register_saved_tensors_hooks(
|
||||
self.pack_hook, self.unpack_hook
|
||||
)
|
||||
|
||||
def __exit__(self, *args: object) -> None:
|
||||
core.eager.reset_saved_tensors_hooks()
|
||||
Reference in New Issue
Block a user