From d287772e6211ace57734d2c86ec0560fc61d640a Mon Sep 17 00:00:00 2001 From: manfredss Date: Fri, 24 Jul 2026 12:10:34 +0800 Subject: [PATCH] try 24 api align --- python/paddle/compat/__init__.py | 564 +++++++++++++++--- python/paddle/compat/nn/__init__.py | 43 +- .../paddle/compat/nn/functional/__init__.py | 166 +++++- python/paddle/compat/nn/functional/sdpa.py | 110 ++++ python/paddle/tensor/compat_softmax.py | 24 +- test/legacy_test/test_compat_minmax.py | 18 +- test/legacy_test/test_compat_split.py | 15 +- 7 files changed, 791 insertions(+), 149 deletions(-) diff --git a/python/paddle/compat/__init__.py b/python/paddle/compat/__init__.py index 752e927d111a3..1d651a3a5512a 100644 --- a/python/paddle/compat/__init__.py +++ b/python/paddle/compat/__init__.py @@ -15,6 +15,7 @@ from __future__ import annotations import importlib +import operator from typing import TYPE_CHECKING, Any, Literal, NamedTuple from typing_extensions import overload @@ -114,6 +115,41 @@ def allclose( >>> print(result2) False """ + if input.dtype != other.dtype: + return paddle.allclose( + input, + other, + rtol=rtol, + atol=atol, + equal_nan=equal_nan, + name=name, + ).item() + + if input.dtype in ( + paddle.float16, + paddle.bfloat16, + paddle.uint8, + paddle.int8, + paddle.int16, + ): + input = input.cast(paddle.float32) + other = other.cast(paddle.float32) + + if input.shape != other.shape: + input, other = paddle.broadcast_tensors([input, other]) + + if input.dtype in (paddle.complex64, paddle.complex128): + close = paddle.logical_or( + paddle.equal(input, other), + paddle.abs(input - other) <= atol + rtol * paddle.abs(other), + ) + if equal_nan: + close = paddle.logical_or( + close, + paddle.logical_and(paddle.isnan(input), paddle.isnan(other)), + ) + return paddle.all(close).item() + return paddle.allclose( input, other, rtol=rtol, atol=atol, equal_nan=equal_nan, name=name ).item() @@ -153,8 +189,8 @@ def equal( >>> print(result1) False """ - if input.dtype == other.dtype: - return paddle.equal_all(input, other).item() + if input.shape != other.shape: + return False common_dtype = promote_types(input.dtype, other.dtype) if input.dtype != common_dtype: @@ -162,7 +198,7 @@ def equal( if other.dtype != common_dtype: other = other.cast(common_dtype) - return paddle.equal_all(input, other).item() + return paddle.all(paddle.equal(input, other)).item() class MedianRetType(NamedTuple): @@ -170,6 +206,109 @@ class MedianRetType(NamedTuple): indices: Tensor +class _ReplaceForwardValue(paddle.autograd.PyLayer): + @staticmethod + def forward(ctx, source, value): + return value.clone() + + @staticmethod + def backward(ctx, grad): + return grad, None + + +def _replace_forward_value(source: Tensor, value: Tensor) -> Tensor: + if not in_dynamic_mode() or source.stop_gradient: + return value + return _ReplaceForwardValue.apply(source, value.detach()) + + +def _cast_reduction_input(input: Tensor, include_bool: bool = False) -> Tensor: + if input.dtype in (paddle.float16, paddle.bfloat16): + return input.cast(paddle.float32) + integral_dtypes = (paddle.uint8, paddle.int8, paddle.int16) + if include_bool: + integral_dtypes += (paddle.bool,) + if input.dtype in integral_dtypes: + return input.cast(paddle.int32) + return input + + +def _normalize_reduction_dim(dim: Any) -> int: + if isinstance(dim, Variable): + if dim.ndim != 0 or dim.dtype not in ( + paddle.uint8, + paddle.int8, + paddle.int16, + paddle.int32, + paddle.int64, + ): + raise TypeError("dim must be a 0-D integral tensor") + if type(dim) is bool: + raise TypeError("dim must be an int, but got bool") + try: + return operator.index(dim) + except TypeError: + raise TypeError( + f"dim must be an int, but got {type(dim).__name__}" + ) from None + + +def _median_with_index( + input: Tensor, dim: int, keepdim: bool, ignore_nan: bool = False +) -> MedianRetType: + work_input = _cast_reduction_input(input) + if input.ndim == 0: + values = input.clone() + return MedianRetType( + values=values, + indices=paddle.zeros([], dtype=paddle.int64, device=input.place), + ) + + axis = dim + input.ndim if dim < 0 else dim + if in_dynamic_mode() and input.shape[axis] == 0: + raise IndexError( + f"median(): Expected reduction dim {dim} to have non-zero size." + ) + + _, sorted_indices = _C_ops.argsort(work_input, axis, False, True) + is_floating = work_input.dtype in (paddle.float32, paddle.float64) + if ignore_nan and is_floating: + valid_count = paddle.sum( + paddle.logical_not(paddle.isnan(work_input)).cast(paddle.int64), + axis=axis, + keepdim=True, + ) + offset = paddle.floor_divide( + paddle.maximum(valid_count - 1, paddle.zeros_like(valid_count)), + paddle.full_like(valid_count, 2), + ) + indices = paddle.take_along_axis(sorted_indices, offset, axis=axis) + else: + offset = (input.shape[axis] - 1) // 2 + indices = paddle.slice( + sorted_indices, + axes=[axis], + starts=[offset], + ends=[offset + 1], + ) + if is_floating and not ignore_nan: + nan_mask = paddle.isnan(work_input) + indices = paddle.where( + paddle.any(nan_mask, axis=axis, keepdim=True), + paddle.argmax(nan_mask.cast(paddle.int32), axis=axis, keepdim=True), + indices, + ) + + indices.stop_gradient = True + values = paddle.take_along_axis(work_input, indices, axis=axis) + if values.dtype != input.dtype: + values = values.cast(input.dtype) + if not keepdim: + values = paddle.squeeze(values, axis=axis) + indices = paddle.squeeze(indices, axis=axis) + return MedianRetType(values=values, indices=indices) + + @ForbidKeywordsDecorator( illegal_keys={"x", "axis"}, func_name="paddle.compat.median", @@ -222,17 +361,41 @@ def median( """ if dim is None: _check_out_status(out, False) - result = paddle.median(input, axis=dim, keepdim=keepdim, mode='min') + work_input = _cast_reduction_input(input) + result = paddle.median( + work_input, axis=dim, keepdim=keepdim, mode='min' + ) + if in_dynamic_mode(): + flat_input = paddle.reshape(work_input, [-1]) + flat_result = paddle.reshape(result, [-1])[0] + matches = flat_input == flat_result + if work_input.dtype in (paddle.float32, paddle.float64): + matches = paddle.logical_or( + matches, + paddle.logical_and( + paddle.isnan(flat_input), paddle.isnan(flat_result) + ), + ) + selected = flat_input[ + paddle.argmax(matches.cast(paddle.int32)) + ].reshape(paddle.shape(result)) + result = _replace_forward_value(result, selected) + if result.dtype != input.dtype: + result = result.cast(input.dtype) if out is not None: paddle.assign(result, out) return out return result else: _check_out_status(out, True) - values, indices = paddle.median( - input, axis=dim, keepdim=keepdim, mode='min', out=out + values, indices = _median_with_index( + input, + _normalize_reduction_dim(dim), + keepdim, ) if out is not None: + paddle.assign(values, out[0]) + paddle.assign(indices, out[1]) return MedianRetType(values=out[0], indices=out[1]) return MedianRetType(values=values, indices=indices) @@ -290,19 +453,24 @@ def nanmedian( """ if dim is None: _check_out_status(out, False) - result = paddle.nanmedian(input, axis=dim, keepdim=keepdim, mode='min') + work_input = _cast_reduction_input(input) + result = paddle.nanmedian( + work_input, axis=dim, keepdim=keepdim, mode='min' + ) + if result.dtype != input.dtype: + result = result.cast(input.dtype) if out is not None: paddle.assign(result, out) return out return result else: _check_out_status(out, True) - values, indices = paddle.nanmedian( - input, axis=dim, keepdim=keepdim, mode='min' + values, indices = _median_with_index( + input, + _normalize_reduction_dim(dim), + keepdim, + ignore_nan=True, ) - # This conversion is needed because PyTorch returns index 0 for all-nan rows, - # while PaddlePaddle returns index -1 for all-nan rows - indices = paddle.maximum(indices, paddle.zeros_like(indices)) if out is not None: paddle.assign(values, out[0]) @@ -360,6 +528,7 @@ def try_get_keys(key): num_args = len(args) total_arg_num = num_args + len(kwargs) + dim_overload = total_arg_num == 2 or "dim" in kwargs or "keepdim" in kwargs if total_arg_num > 2: raise invalid_arguments_exception() elif total_arg_num == 2: @@ -371,9 +540,7 @@ def try_get_keys(key): else: dim_or_other = try_get_keys("dim") keepdim = try_get_keys("keepdim") - if dim_or_other is None or isinstance( - dim_or_other, (Variable, paddle.pir.Value) - ): + if dim_or_other is None: raise invalid_arguments_exception() elif total_arg_num == 1: if num_args: @@ -388,14 +555,14 @@ def try_get_keys(key): if dim_or_other is None: raise invalid_arguments_exception() - if ( - dim_or_other is not None - and not isinstance(dim_or_other, (Variable, paddle.pir.Value)) - and type(dim_or_other) is not int - ): - raise invalid_arguments_exception( - f"The second input must be int or Tensor or implicit None in compat.{func_name}, but received {type(dim_or_other)}.\n" - ) + is_tensor = isinstance(dim_or_other, (Variable, paddle.pir.Value)) + if dim_or_other is not None and not dim_overload and not is_tensor: + dim_overload = True + if dim_overload: + try: + dim_or_other = _normalize_reduction_dim(dim_or_other) + except TypeError as error: + raise invalid_arguments_exception(str(error)) from None return dim_or_other, keepdim @@ -404,7 +571,9 @@ def _min_max_tensor_allow_grad(input: Tensor): """Prevent integral input tensor type to have `stop_gradient=False`""" in_dtype = input.dtype if ( - in_dtype == paddle.int32 + in_dtype == paddle.bool + or in_dtype == paddle.int8 + or in_dtype == paddle.int32 or in_dtype == paddle.int64 or in_dtype == paddle.uint8 or in_dtype == paddle.int16 @@ -415,18 +584,115 @@ def _min_max_tensor_allow_grad(input: Tensor): ) -def _min_max_allow_cpu_composite(input: Tensor): - """paddle.min/argmin(max/argmax), paddle.take_along_axis reject the following types""" - in_dtype = input.dtype - if ( - in_dtype == paddle.float16 - or in_dtype == paddle.bfloat16 - or in_dtype == paddle.int16 - ): - raise TypeError( - f"Non-CUDA GPU placed Tensor does not have '{in_dtype}' op registered.\n" - "Paddle support following DataTypes: int32, int64, float64, float32, uint8" +class _MinMaxReduceAll(paddle.autograd.PyLayer): + @staticmethod + def forward(ctx, input, is_min): + result = paddle.min(input) if is_min else paddle.max(input) + if input.dtype in (paddle.float32, paddle.float64): + zero_mask = input == 0 + if is_min: + signed_zero_mask = paddle.logical_and( + zero_mask, paddle.signbit(input) + ) + else: + signed_zero_mask = paddle.logical_and( + zero_mask, paddle.logical_not(paddle.signbit(input)) + ) + signed_zero = paddle.copysign( + paddle.zeros_like(result), + paddle.full_like(result, -1.0 if is_min else 1.0), + ) + result = paddle.where( + paddle.logical_and(result == 0, paddle.any(signed_zero_mask)), + signed_zero, + result, + ) + ctx.save_for_backward(input, result) + return result + + @staticmethod + def backward(ctx, grad): + input, result = ctx.saved_tensor() + mask = input == result + if input.dtype in (paddle.float32, paddle.float64): + mask = paddle.where(paddle.isnan(result), paddle.isnan(input), mask) + mask = mask.cast(input.dtype) + return grad * mask / paddle.sum(mask) + + +def _min_max_reduce_all(input: Tensor, is_min: bool) -> Tensor: + work_input = _cast_reduction_input(input, include_bool=True) + if in_dynamic_mode(): + result = _MinMaxReduceAll.apply(work_input, is_min) + else: + result = paddle.min(work_input) if is_min else paddle.max(work_input) + return result.cast(input.dtype) if result.dtype != input.dtype else result + + +def _min_max_reduce_dim_fallback( + input: Tensor, dim: int, keepdim: bool, is_min: bool +) -> MinMaxRetType: + axis = dim + input.ndim if dim < 0 else dim + if input.shape[axis] == 0: + raise IndexError( + f"{'min' if is_min else 'max'}(): Expected reduction dim " + f"{dim} to have non-zero size." + ) + + work_input = _cast_reduction_input(input, include_bool=True) + index_op = paddle.argmin if is_min else paddle.argmax + indices = index_op(work_input, axis=axis, keepdim=True) + if work_input.dtype in (paddle.float32, paddle.float64): + nan_mask = paddle.isnan(work_input) + indices = paddle.where( + paddle.any(nan_mask, axis=axis, keepdim=True), + paddle.argmax(nan_mask.cast(paddle.int32), axis=axis, keepdim=True), + indices, + ) + indices.stop_gradient = True + values = paddle.take_along_axis(work_input, indices, axis=axis) + if values.dtype != input.dtype: + values = values.cast(input.dtype) + if not keepdim: + values = paddle.squeeze(values, axis=axis) + indices = paddle.squeeze(indices, axis=axis) + return MinMaxRetType(values=values, indices=indices) + + +def _min_max_elementwise(input: Tensor, other: Tensor, is_min: bool) -> Tensor: + common_dtype = promote_types(input.dtype, other.dtype) + if input.dtype != common_dtype: + input = input.cast(common_dtype) + if other.dtype != common_dtype: + other = other.cast(common_dtype) + + work_input = _cast_reduction_input(input, include_bool=True) + work_other = ( + other.cast(work_input.dtype) + if other.dtype != work_input.dtype + else other + ) + op = _C_ops.minimum if is_min else _C_ops.maximum + result = op(work_input, work_other) + if work_input.dtype in (paddle.float32, paddle.float64): + nan_mask = paddle.logical_or( + paddle.isnan(work_input), paddle.isnan(work_other) ) + result = paddle.where(nan_mask, work_input + work_other, result) + if in_dynamic_mode(): + broadcast_input, broadcast_other = paddle.broadcast_tensors( + [work_input, work_other] + ) + zero_tie = paddle.logical_and( + broadcast_input == 0, broadcast_other == 0 + ) + signed_zero = paddle.copysign( + paddle.zeros_like(result), broadcast_input + ) + result = _replace_forward_value( + result, paddle.where(zero_tie, signed_zero, result) + ) + return result.cast(common_dtype) if result.dtype != common_dtype else result @ForbidKeywordsDecorator( @@ -548,25 +814,14 @@ def min( if dim_or_other is None: # paddle.min and paddle.amin actually shares the same grad op (ReduceAminKernel) _check_out_status(out, False) - ret = paddle.min(input) + ret = _min_max_reduce_all(input, is_min=True) elif isinstance(dim_or_other, int): _check_out_status(out, True) if input.ndim: if in_dynamic_mode() and not input.place.is_gpu_place(): - _min_max_allow_cpu_composite(input) - # CPUPlace and other placements are implemented by composition - - indices = paddle.argmin(input, axis=dim_or_other, keepdim=True) - values = paddle.take_along_axis( - input, indices, axis=dim_or_other + ret = _min_max_reduce_dim_fallback( + input, dim_or_other, keepdim, is_min=True ) - if keepdim: - ret = MinMaxRetType(values=values, indices=indices) - else: - ret = MinMaxRetType( - values=values.squeeze_(axis=dim_or_other), - indices=indices.squeeze_(axis=dim_or_other), - ) else: vals, inds = _C_ops.min_with_index( input, dim_or_other, keepdim, False @@ -575,21 +830,23 @@ def min( ret = MinMaxRetType(values=vals, indices=inds) else: ret = MinMaxRetType( - values=input, + values=input.clone(), indices=paddle.zeros( [], dtype=paddle.int64, device=input.place ), ) else: _check_out_status(out, False) - ret = _C_ops.minimum(input, dim_or_other) + ret = _min_max_elementwise(input, dim_or_other, is_min=True) if out is not None: if isinstance(ret, MinMaxRetType): paddle.assign(ret.values, out[0]) paddle.assign(ret.indices, out[1]) + return MinMaxRetType(values=out[0], indices=out[1]) else: paddle.assign(ret, out) + return out return ret @@ -711,23 +968,14 @@ def max( ret = None if dim_or_other is None: _check_out_status(out, False) - ret = paddle.max(input) + ret = _min_max_reduce_all(input, is_min=False) elif isinstance(dim_or_other, int): _check_out_status(out, True) if input.ndim: if in_dynamic_mode() and not input.place.is_gpu_place(): - _min_max_allow_cpu_composite(input) - indices = paddle.argmax(input, axis=dim_or_other, keepdim=True) - values = paddle.take_along_axis( - input, indices, axis=dim_or_other + ret = _min_max_reduce_dim_fallback( + input, dim_or_other, keepdim, is_min=False ) - if keepdim: - ret = MinMaxRetType(values=values, indices=indices) - else: - ret = MinMaxRetType( - values=values.squeeze_(axis=dim_or_other), - indices=indices.squeeze_(axis=dim_or_other), - ) else: vals, inds = _C_ops.max_with_index( input, dim_or_other, keepdim, False @@ -736,21 +984,23 @@ def max( ret = MinMaxRetType(values=vals, indices=inds) else: ret = MinMaxRetType( - values=input, + values=input.clone(), indices=paddle.zeros( [], dtype=paddle.int64, device=input.place ), ) else: _check_out_status(out, False) - ret = _C_ops.maximum(input, dim_or_other) + ret = _min_max_elementwise(input, dim_or_other, is_min=False) if out is not None: if isinstance(ret, MinMaxRetType): paddle.assign(ret.values, out[0]) paddle.assign(ret.indices, out[1]) + return MinMaxRetType(values=out[0], indices=out[1]) else: paddle.assign(ret, out) + return out return ret @@ -881,13 +1131,35 @@ def sort( [1, 0, 3, 2]])) """ _check_out_status(out, expect_multiple=True) - outputs, indices = _C_ops.argsort(input, dim, descending, stable) + work_input = ( + input.cast(paddle.int16) + if input.dtype in (paddle.bool, paddle.int8) + else input + ) + outputs, indices = _C_ops.argsort(work_input, dim, descending, stable) + indices.stop_gradient = True + if outputs.dtype != input.dtype: + outputs = outputs.cast(input.dtype) if out is not None: paddle.assign(outputs, out[0]) paddle.assign(indices, out[1]) + return SortRetType(values=out[0], indices=out[1]) return SortRetType(values=outputs, indices=indices) +class _UniqueBackwardNotImplemented(paddle.autograd.PyLayer): + @staticmethod + def forward(ctx, input, output, op_name): + ctx.op_name = op_name + return output.clone() + + @staticmethod + def backward(ctx, grad): + raise NotImplementedError( + f"the derivative for '{ctx.op_name}' is not implemented." + ) + + @overload def unique( input: Tensor, @@ -988,14 +1260,74 @@ def unique( [[2, 1, 3], [3, 0, 1]]) """ - return paddle.unique( - input, + work_input = _cast_reduction_input(input, include_bool=True) + result = paddle.unique( + work_input, return_inverse=return_inverse, return_counts=return_counts, axis=dim, sorted=sorted, ) + if return_inverse and return_counts: + output, inverse, counts = result + elif return_inverse: + output, inverse = result + counts = None + elif return_counts: + output, counts = result + inverse = None + else: + output = result + inverse = counts = None + + if dim is None and inverse is not None: + inverse = inverse.reshape(input.shape) + + if ( + dim is None + and (inverse is not None or counts is not None) + and work_input.dtype in (paddle.float32, paddle.float64) + ): + flat_input = work_input.reshape([-1]) + nan_mask = paddle.isnan(flat_input) + nan_count = paddle.sum(nan_mask.cast(paddle.int64)) + non_nan_count = paddle.shape(output)[0] - nan_count + if inverse is not None: + nan_inverse = ( + non_nan_count + paddle.cumsum(nan_mask.cast(paddle.int64)) - 1 + ) + inverse = paddle.where(nan_mask, nan_inverse, inverse.reshape([-1])) + inverse = inverse.reshape(input.shape) + if counts is not None: + counts = paddle.where( + paddle.arange(paddle.shape(output)[0]) >= non_nan_count, + paddle.ones_like(counts), + counts, + ) + + if output.dtype != input.dtype: + output = output.cast(input.dtype) + if ( + in_dynamic_mode() + and not input.stop_gradient + and input.dtype + in (paddle.float16, paddle.float32, paddle.float64, paddle.bfloat16) + ): + output = _UniqueBackwardNotImplemented.apply( + input, + output.detach(), + 'unique_dim' if dim is not None else '_unique2', + ) + + if return_inverse and return_counts: + return output, inverse, counts + if return_inverse: + return output, inverse + if return_counts: + return output, counts + return output + @ForbidKeywordsDecorator( illegal_keys={"x", "num_or_sections", "axis", "name"}, @@ -1089,48 +1421,84 @@ def GetShapeOnDimInRange(shape, dim: int) -> int: shape_val = int(section_size.item(0)) else: shape_val = section_size - if section_size < 0: + if shape_val < 0: raise ValueError( f"paddle.compat.split expects split_sizes have only non-negative entries, but got size = {section_size} on dim {i}" ) if in_dynamic_mode(): - if isinstance(dim, Variable): - dim = dim.item(0) - assert dim + len(tensor.shape) >= 0, "(rank(x) + dim) must >= 0" - dim = (dim + len(tensor.shape)) if dim < 0 else dim + integral_dtypes = ( + paddle.uint8, + paddle.int8, + paddle.int16, + paddle.int32, + paddle.int64, + ) + + def to_index(value, argument_name): + if isinstance(value, Variable): + if ( + int(value.numel().item()) != 1 + or value.dtype not in integral_dtypes + ): + raise TypeError(f"{argument_name} must be an integer") + return int(value.item()) + if type(value) is bool: + raise TypeError(f"{argument_name} must be an integer") + try: + return operator.index(value) + except TypeError: + raise TypeError(f"{argument_name} must be an integer") from None + + if tensor.ndim == 0: + raise RuntimeError("split expects at least a 1-dimensional tensor") + if isinstance(dim, Variable) and dim.ndim != 0: + raise TypeError("dim tensor must be a 0-D integral tensor") + dim = to_index(dim, "dim") + shape_on_dim = GetShapeOnDimInRange(tensor.shape, dim) + dim = dim + tensor.ndim if dim < 0 else dim if isinstance(split_size_or_sections, (list, tuple)): - if paddle.utils._contain_var(split_size_or_sections): - for index, item in enumerate(split_size_or_sections): - if isinstance(item, Variable): - split_size_or_sections[index] = split_size_or_sections[ - index - ].item() - elif not isinstance(split_size_or_sections, int): - raise TypeError( - "The type of 'split_size_or_sections' in split must be int, list or tuple in imperative mode, but " - f"received {type(split_size_or_sections)}." + split_size_or_sections = [ + to_index(item, "split_sizes") for item in split_size_or_sections + ] + else: + split_size_or_sections = to_index( + split_size_or_sections, "split_size_or_sections" ) if isinstance(split_size_or_sections, int): - # check whether shape is divisible - assert split_size_or_sections > 0, ( - 'split_size_or_sections must be greater than 0.' - ) - - split_size_or_sections = GetSplitSize( - split_size_or_sections, GetShapeOnDimInRange(tensor.shape, dim) - ) - - if isinstance(split_size_or_sections, list): - return tuple(_C_ops.split(tensor, split_size_or_sections, dim)) + if split_size_or_sections < 0 or ( + split_size_or_sections == 0 and shape_on_dim != 0 + ): + raise RuntimeError( + "split_size can only be 0 if dimension size is 0" + ) + if shape_on_dim == 0: + split_size_or_sections = [0] else: - return tuple( - _C_ops.split_with_num(tensor, split_size_or_sections, dim) + num_complete_sections, remaining_num = divmod( + shape_on_dim, split_size_or_sections ) - else: - return tuple(_C_ops.split(tensor, split_size_or_sections, dim)) + split_size_or_sections = [ + split_size_or_sections for _ in range(num_complete_sections) + ] + if remaining_num: + split_size_or_sections.append(remaining_num) + elif sum(split_size_or_sections) != shape_on_dim: + raise RuntimeError( + "split_sizes must sum exactly to the selected dimension " + f"size {shape_on_dim}, but got {split_size_or_sections}" + ) + + outputs = [] + start = 0 + for section_size in split_size_or_sections: + slices = [slice(None)] * tensor.ndim + slices[dim] = slice(start, start + section_size) + outputs.append(tensor[tuple(slices)]) + start += section_size + return tuple(outputs) else: if isinstance(dim, paddle.pir.Value): raise TypeError( diff --git a/python/paddle/compat/nn/__init__.py b/python/paddle/compat/nn/__init__.py index e401845d96c7b..a443791744149 100644 --- a/python/paddle/compat/nn/__init__.py +++ b/python/paddle/compat/nn/__init__.py @@ -532,19 +532,24 @@ def __init__( stride: Size2 = 1, ) -> None: super().__init__(kernel_size, dilation, padding, stride) + self.kernel_size = kernel_size + self.dilation = dilation + self.padding = padding + self.stride = stride def forward(self, input: Tensor) -> Tensor: - def to_list_if_necessary(x): - if isinstance(x, (paddle.pir.Value, paddle.Tensor)): - x = x.tolist() - return x - - return nn.functional.unfold( + return functional._unfold( input, - kernel_sizes=to_list_if_necessary(self.kernel_sizes), - strides=to_list_if_necessary(self.strides), - paddings=to_list_if_necessary(self.paddings), - dilations=to_list_if_necessary(self.dilations), + self.kernel_size, + self.dilation, + self.padding, + self.stride, + ) + + def extra_repr(self) -> str: + return ( + f"kernel_size={self.kernel_size}, dilation={self.dilation}, " + f"padding={self.padding}, stride={self.stride}" ) @@ -646,12 +651,24 @@ def __init__( ) self.in_features = in_features self.out_features = out_features + default_initializer = ( + nn.initializer.Constant(0.0) + if self._dtype + in ( + 'complex64', + 'complex128', + paddle.complex64, + paddle.complex128, + ) + else None + ) self.weight = self.create_parameter( shape=[out_features, in_features], attr=None, dtype=self._dtype, is_bias=False, device=device, + default_initializer=default_initializer, ) self.bias = None if bias: @@ -661,6 +678,7 @@ def __init__( dtype=self._dtype, is_bias=True, device=device, + default_initializer=default_initializer, ) # The same parameter initialization as PyTorch self.reset_parameters() @@ -816,14 +834,15 @@ class Softmax(nn.Layer, metaclass=_CompatClassMeta): ) def __init__(self, dim: int | None = None) -> None: nn.Layer.__init__(self) + self.dim = dim self._dim = dim self._dtype = None def forward(self, input: Tensor) -> Tensor: - return functional.softmax(input, self._dim) + return functional.softmax(input, self.dim) def extra_repr(self) -> str: - return f"dim={self._dim}" + return f"dim={self.dim}" class SmoothL1Loss(nn.Layer, metaclass=_CompatClassMeta): diff --git a/python/paddle/compat/nn/functional/__init__.py b/python/paddle/compat/nn/functional/__init__.py index 5929b0e750f6e..c7d07d8d79be3 100644 --- a/python/paddle/compat/nn/functional/__init__.py +++ b/python/paddle/compat/nn/functional/__init__.py @@ -66,6 +66,52 @@ def _check_valid_pad_len(pad_len, x_dim, is_constant): ) +def _crop_negative_padding(input, pad): + if not isinstance(pad, (list, tuple)) or len(pad) > input.ndim * 2: + return input, pad + + normalized_pad = [] + for value in pad: + if ( + in_dynamic_mode() + and isinstance(value, paddle.Tensor) + and value.numel() == 1 + ): + value = int(value.item()) + normalized_pad.append(value) + + if not any( + isinstance(value, int) and not isinstance(value, bool) and value < 0 + for value in normalized_pad + ): + return input, pad + + slices = [slice(None)] * input.ndim + for pair_idx in range(len(normalized_pad) // 2): + left, right = normalized_pad[2 * pair_idx : 2 * pair_idx + 2] + if not isinstance(left, int) or not isinstance(right, int): + continue + axis = input.ndim - pair_idx - 1 + crop_left = max(-left, 0) + crop_right = max(-right, 0) + if ( + input.shape[axis] >= 0 + and crop_left + crop_right > input.shape[axis] + ): + raise RuntimeError("narrow(): length must be non-negative.") + slices[axis] = slice( + crop_left, + None if crop_right == 0 else -crop_right, + ) + normalized_pad[2 * pair_idx] = max(left, 0) + normalized_pad[2 * pair_idx + 1] = max(right, 0) + + input = input[tuple(slices)] + if isinstance(pad, tuple): + normalized_pad = tuple(normalized_pad) + return input, normalized_pad + + @ForbidKeywordsDecorator( illegal_keys={"x", "name", "data_format", "pad_from_left_axis"}, func_name="paddle.compat.nn.functional.pad", @@ -75,7 +121,7 @@ def pad( input: Tensor, pad: ShapeLike, mode: _PaddingTensorMode = 'constant', - value: float = 0.0, + value: float | None = None, ) -> Tensor: """ @@ -138,6 +184,31 @@ def pad( if isinstance(pad, (Variable, paddle.Tensor)) and pad.size == 0: return input.clone() + if value is None: + value = 0.0 + + if mode in ("constant", "circular"): + input, pad = _crop_negative_padding(input, pad) + original_dtype = input.dtype + fallback_dtype = None + if in_dynamic_mode(): + if input.dtype in (paddle.uint8, paddle.int8, paddle.int16) or ( + input.dtype == paddle.bool and mode in ("constant", "circular") + ): + fallback_dtype = paddle.int32 + elif not input.place.is_gpu_place() and input.dtype in ( + paddle.float16, + paddle.bfloat16, + ): + fallback_dtype = paddle.float32 + if fallback_dtype is not None: + if original_dtype == paddle.bool and mode == "constant": + value = float(bool(value)) + input = input.cast(fallback_dtype) + + def restore_dtype(out): + return out.cast(original_dtype) if fallback_dtype is not None else out + if ( mode == "constant" and isinstance(pad, (list, tuple)) @@ -164,7 +235,7 @@ def pad( if isinstance(pad_value, paddle.pir.Value) else float(pad_value) ) - return _C_ops.pad(input, paddings, pad_val) + return restore_dtype(_C_ops.pad(input, paddings, pad_val)) assert x_dim >= 1 and x_dim <= 5, ( f"Input tensor dimension must be in [1-5] but got {x_dim}" @@ -204,8 +275,8 @@ def pad( "NCDHW", ) if ndim_to_unsqueeze: - return out.squeeze(axis=ndim_to_unsqueeze) - return out + out = out.squeeze(axis=ndim_to_unsqueeze) + return restore_dtype(out) @ForbidKeywordsDecorator( @@ -276,10 +347,39 @@ def linear(input: Tensor, weight: Tensor, bias: Tensor | None = None) -> Tensor: [3.50000000, 3.50000000, 3.50000000, 3.50000000], [5.50000000, 5.50000000, 5.50000000, 5.50000000]]) """ + if ( + in_dynamic_mode() + and input.dtype == weight.dtype + and (bias is None or input.dtype == bias.dtype) + and ( + input.dtype in (paddle.uint8, paddle.int8, paddle.int16) + or ( + not input.place.is_gpu_place() + and input.dtype in (paddle.float16, paddle.bfloat16) + ) + ) + ): + output_dtype = input.dtype + compute_dtype = ( + paddle.float32 + if input.dtype in (paddle.float16, paddle.bfloat16) + else paddle.int32 + ) + out = _C_ops.matmul( + input.cast(compute_dtype), + weight.cast(compute_dtype), + False, + True, + ) + if bias is not None: + out = _C_ops.add(out, bias.cast(compute_dtype)) + return out.cast(output_dtype) + if ( paddle.get_flags("FLAGS_use_legacy_linear")["FLAGS_use_legacy_linear"] or not paddle.is_compiled_with_cuda() or not paddle.framework.in_dynamic_or_pir_mode() + or input.dtype in (paddle.complex64, paddle.complex128) ): # Fallback to old logic when in non-cuda or legacy mode. out = _C_ops.matmul(input, weight, False, True) @@ -295,6 +395,51 @@ def linear(input: Tensor, weight: Tensor, bias: Tensor | None = None) -> Tensor: return _C_ops.matmul(input, weight.contiguous(), False, True) +def _unfold( + input: Tensor, + kernel_size: Size2, + dilation: Size2 = 1, + padding: Size2 = 0, + stride: Size2 = 1, +) -> Tensor: + def to_list_if_necessary(x): + if isinstance(x, (paddle.pir.Value, paddle.Tensor)): + return x.tolist() + return x + + def native_unfold(x): + return paddle.nn.functional.unfold( + x=x, + kernel_sizes=to_list_if_necessary(kernel_size), + strides=to_list_if_necessary(stride), + paddings=to_list_if_necessary(padding), + dilations=to_list_if_necessary(dilation), + ) + + is_unbatched = input.ndim == 3 + if is_unbatched: + input = input.unsqueeze(0) + + if in_dynamic_mode(): + output_dtype = input.dtype + if output_dtype in (paddle.complex64, paddle.complex128): + out = paddle.complex( + native_unfold(paddle.real(input)), + native_unfold(paddle.imag(input)), + ) + elif output_dtype == paddle.bool or ( + not input.place.is_gpu_place() + and output_dtype in (paddle.float16, paddle.bfloat16) + ): + out = native_unfold(input.cast(paddle.float32)).cast(output_dtype) + else: + out = native_unfold(input) + else: + out = native_unfold(input) + + return out.squeeze(0) if is_unbatched else out + + @ForbidKeywordsDecorator( illegal_keys={ "x", @@ -373,18 +518,7 @@ def unfold( >>> y = F.unfold(x, [3, 3], 1, 1, 1) """ - def to_list_if_necessary(x): - if isinstance(x, (paddle.pir.Value, paddle.Tensor)): - x = x.tolist() - return x - - return paddle.nn.functional.unfold( - x=input, - kernel_sizes=to_list_if_necessary(kernel_size), - strides=to_list_if_necessary(stride), - paddings=to_list_if_necessary(padding), - dilations=to_list_if_necessary(dilation), - ) + return _unfold(input, kernel_size, dilation, padding, stride) @ForbidKeywordsDecorator( diff --git a/python/paddle/compat/nn/functional/sdpa.py b/python/paddle/compat/nn/functional/sdpa.py index 9df687c84ffda..050ff262eb30a 100644 --- a/python/paddle/compat/nn/functional/sdpa.py +++ b/python/paddle/compat/nn/functional/sdpa.py @@ -16,12 +16,97 @@ from typing import TYPE_CHECKING +import paddle import paddle.nn.functional as F if TYPE_CHECKING: from paddle import Tensor +def _math_scaled_dot_product_attention( + query, + key, + value, + attn_mask, + dropout_p, + is_causal, + scale, + enable_gqa, +): + output_dtype = query.dtype + query_place = getattr(query, "place", None) + if ( + query_place is not None + and not query_place.is_gpu_place() + and query.dtype in (paddle.float16, paddle.bfloat16) + ): + query = query.cast(paddle.float32) + key = key.cast(paddle.float32) + value = value.cast(paddle.float32) + if attn_mask is not None and attn_mask.dtype != paddle.bool: + attn_mask = attn_mask.cast(paddle.float32) + + if enable_gqa: + query_heads = query.shape[-3] + key_heads = key.shape[-3] + value_heads = value.shape[-3] + if ( + key_heads == 0 + or value_heads == 0 + or query_heads % key_heads != 0 + or query_heads % value_heads != 0 + ): + raise ValueError( + "The number of query heads must be divisible by the number " + "of key/value heads when enable_gqa=True." + ) + key_repeats = query_heads // key_heads + value_repeats = query_heads // value_heads + if key_repeats != 1: + key = paddle.repeat_interleave(key, key_repeats, axis=-3) + if value_repeats != 1: + value = paddle.repeat_interleave(value, value_repeats, axis=-3) + + scale_factor = query.shape[-1] ** -0.5 if scale is None else scale + scores = paddle.matmul(query, key, transpose_y=True) * scale_factor + + if is_causal: + causal_mask = paddle.ones( + [query.shape[-2], key.shape[-2]], dtype=paddle.bool + ).tril() + scores = paddle.where( + causal_mask, + scores, + paddle.full_like(scores, -float("inf")), + ) + if attn_mask is not None: + if attn_mask.dtype == paddle.bool: + scores = paddle.where( + attn_mask, + scores, + paddle.full_like(scores, -float("inf")), + ) + else: + scores = scores + attn_mask + + has_unmasked_score = paddle.any( + scores != -float("inf"), axis=-1, keepdim=True + ) + safe_scores = paddle.where( + has_unmasked_score, scores, paddle.zeros_like(scores) + ) + weights = F.softmax(safe_scores, axis=-1) + weights = paddle.where( + paddle.logical_and(has_unmasked_score, scores != -float("inf")), + weights, + paddle.zeros_like(weights), + ) + if dropout_p > 0.0: + weights = F.dropout(weights, p=dropout_p, training=True) + out = paddle.matmul(weights, value) + return out if out.dtype == output_dtype else out.cast(output_dtype) + + def scaled_dot_product_attention( query: Tensor, key: Tensor, @@ -105,6 +190,31 @@ def scaled_dot_product_attention( "Explicit attn_mask should not be set when is_causal=True" ) + query_place = getattr(query, "place", None) + use_math_fallback = ( + query.ndim not in (3, 4) + or scale == 0 + or ( + query_place is not None + and not query_place.is_gpu_place() + and ( + query.dtype in (paddle.float16, paddle.bfloat16) + or attn_mask is not None + ) + ) + ) + if use_math_fallback: + return _math_scaled_dot_product_attention( + query, + key, + value, + attn_mask, + dropout_p, + is_causal, + scale, + enable_gqa, + ) + query, key, value = ( query.swapaxes(-3, -2), key.swapaxes(-3, -2), diff --git a/python/paddle/tensor/compat_softmax.py b/python/paddle/tensor/compat_softmax.py index 2afefbb53a58d..e80c79eb17f50 100644 --- a/python/paddle/tensor/compat_softmax.py +++ b/python/paddle/tensor/compat_softmax.py @@ -27,6 +27,26 @@ from paddle._typing import DTypeLike +def _softmax_fallback(op, input, dim, out): + place = getattr(input, "place", None) + if ( + place is not None + and not place.is_gpu_place() + and input.dtype + in ( + core.DataType.FLOAT16, + core.DataType.BFLOAT16, + core.VarDesc.VarType.FP16, + core.VarDesc.VarType.BF16, + ) + ): + output_dtype = input.dtype + result = op(_C_ops.cast(input, core.DataType.FLOAT32), dim) + result = _C_ops.cast(result, output_dtype) + return result if out is None else _C_ops.assign_out_(result, out) + return op(input, dim, out=out) + + @ForbidKeywordsIgnoreOneParamDecorator( illegal_keys={"x", "axis", "name"}, ignore_param=('_stacklevel', 2, int), @@ -177,7 +197,7 @@ def softmax( dtype = convert_nptype_to_datatype_or_vartype(dtype) if in_dynamic_or_pir_mode(): outs_cast = input if dtype is None else _C_ops.cast(input, dtype) - return _C_ops.softmax(outs_cast, dim, out=out) + return _softmax_fallback(_C_ops.softmax, outs_cast, dim, out) @ForbidKeywordsIgnoreOneParamDecorator( @@ -258,4 +278,4 @@ def log_softmax( if in_dynamic_or_pir_mode(): outs_cast = input if dtype is None else _C_ops.cast(input, dtype) - return _C_ops.log_softmax(outs_cast, dim, out=out) + return _softmax_fallback(_C_ops.log_softmax, outs_cast, dim, out) diff --git a/test/legacy_test/test_compat_minmax.py b/test/legacy_test/test_compat_minmax.py index 9212f8a163279..82d17edb36286 100644 --- a/test/legacy_test/test_compat_minmax.py +++ b/test/legacy_test/test_compat_minmax.py @@ -282,10 +282,6 @@ def test_error_handling(self): f"{self.test_op_name}() received unexpected keyword argument 'axis'. " f"\nDid you mean to use {self.origin_op_name}() instead?" ) - err_msg4 = ( - "Non-CUDA GPU placed Tensor does not have 'paddle.float16' op registered.\n" - "Paddle support following DataTypes: int32, int64, float64, float32, uint8" - ) err_msg5 = ( "input should be a tensor, but got an instance with type 'list'" ) @@ -336,9 +332,8 @@ def test_error_handling(self): with self.assertRaises(TypeError) as cm: self.test_op(input_ts, dim=paddle.to_tensor([0])) - # Tensor input for dim case 2 - with self.assertRaises(TypeError) as cm: - self.test_op(input_ts, dim=paddle.to_tensor(0)) + # 0-D integral Tensor is accepted as dim + self.test_op(input_ts, dim=paddle.to_tensor(0)) # Tensor input for dim case 3 with self.assertRaises(TypeError) as cm: @@ -370,12 +365,9 @@ def test_error_handling(self): self.test_op(input_ts, axis=0) self.assertEqual(str(cm.exception), err_msg3) - # Rejected on CPU types - with self.assertRaises(TypeError) as cm: - tensor = paddle.to_tensor([1, 2, 3], dtype="float16") - cpu_tensor = tensor.to("cpu") - self.test_op(cpu_tensor, dim=0) - self.assertEqual(str(cm.exception), err_msg4) + # Supported on CPU through the compatibility fallback + tensor = paddle.to_tensor([1, 2, 3], dtype="float16") + self.test_op(tensor.to("cpu"), dim=0) # Wrong input type with self.assertRaises(TypeError) as cm: diff --git a/test/legacy_test/test_compat_split.py b/test/legacy_test/test_compat_split.py index 6922b58185512..42e3f3fe27831 100644 --- a/test/legacy_test/test_compat_split.py +++ b/test/legacy_test/test_compat_split.py @@ -109,7 +109,7 @@ def test_empty_dim(self): def test_split_with_one_block(self): """Resulting tuple should be of length 1""" in_tensor = paddle.arange(60, dtype=paddle.float32).reshape([3, 4, 5]) - self._compare_with_origin(in_tensor, 5, paddle.to_tensor([-1])) + self._compare_with_origin(in_tensor, 5, paddle.to_tensor(-1)) self._compare_with_origin(in_tensor, [5], paddle.to_tensor(2)) def test_edge_cases(self): @@ -118,13 +118,17 @@ def test_edge_cases(self): s1, s2 = split(x, [3, 2]) np.testing.assert_allclose(s1.numpy(), [0, 1, 2]) np.testing.assert_allclose(s2.numpy(), [3, 4]) + tensor_size_result = split(x, paddle.to_tensor(2)) + self.assertEqual( + [item.shape[0] for item in tensor_size_result], [2, 2, 1] + ) x = paddle.rand([2, 2, 2]) a, b = split(x, 1, 2) self.assertEqual(a.shape, [2, 2, 1]) # invalid split sections - with self.assertRaises(ValueError): + with self.assertRaises(RuntimeError): split(x, [3, 1], 1) # invalid split axis @@ -146,11 +150,7 @@ def test_error_hint(self): msg_gt_3 = "(InvalidArgument) The dim is expected to be in range of [-3, 3), but got 3" msg_gt_4 = "paddle.compat.split expects split_sizes have only non-negative entries, but got size = -5 on dim 2" - split_size = paddle.to_tensor([3]) - msg_gt_5 = ( - "The type of 'split_size_or_sections' in split must be int, list or tuple in imperative mode, but " - f"received {type(split_size)}." - ) + split_size = paddle.to_tensor([3, 3]) with self.assertRaises(TypeError) as cm: tensors = paddle.split(tensor=x, split_size_or_sections=3, dim=0) @@ -170,7 +170,6 @@ def test_error_hint(self): with self.assertRaises(TypeError) as cm: tensors = split(x, split_size, 1) - self.assertEqual(str(cm.exception), msg_gt_5) class TestFunctionalSplit(unittest.TestCase):