Files
paddlepaddle--paddle/paddle/phi/infermeta/backward.cc
T
2026-07-13 12:40:42 +08:00

2443 lines
86 KiB
C++

/* 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. */
#include "paddle/phi/infermeta/backward.h"
#include "paddle/phi/common/type_traits.h"
#include "paddle/phi/core/utils/data_type.h"
#include "paddle/phi/kernels/funcs/axis_utils.h"
namespace phi {
void AffineGridGradInferMeta(const MetaTensor& output_grad,
const IntArray& outputShape,
bool align_corners,
MetaTensor* input_grad) {
if (input_grad) {
auto output_dims = output_grad.dims();
if (output_dims.size() == 4) {
input_grad->set_dims(make_ddim({output_dims[0], 2, 3}));
} else {
input_grad->set_dims(make_ddim({output_dims[0], 3, 4}));
}
}
}
void AngleGradInferMeta(const MetaTensor& x,
const MetaTensor& out_grad,
MetaTensor* x_grad) {
UnchangedInferMeta(x, x_grad);
}
void BatchFCGradInferMeta(const MetaTensor& input,
const MetaTensor& w,
const MetaTensor& bias,
const MetaTensor& out_grad,
MetaTensor* input_grad,
MetaTensor* w_grad,
MetaTensor* bias_grad) {
input_grad->set_dims(input.dims());
input_grad->set_dtype(input.dtype());
w_grad->set_dims(w.dims());
w_grad->set_dtype(w.dtype());
bias_grad->set_dims(bias.dims());
bias_grad->set_dtype(bias.dtype());
}
void BilinearGradInferMeta(const MetaTensor& x,
const MetaTensor& y,
const MetaTensor& weight,
const MetaTensor& dout,
MetaTensor* dx,
MetaTensor* dy,
MetaTensor* dweight,
MetaTensor* dbias) {
auto x_dims = x.dims();
auto y_dims = y.dims();
auto weight_dims = weight.dims();
auto out_dims = dout.dims();
PADDLE_ENFORCE_EQ(
out_dims.size(),
2UL,
errors::InvalidArgument("The input(Out@GRAD) must be a 2D Tensor."));
PADDLE_ENFORCE_EQ(
x_dims[0],
out_dims[0],
errors::InvalidArgument(
"The first dimension(batch_size) of input(Out@GRAD) must be "
"equal to the first dimension of the Input(X)."));
PADDLE_ENFORCE_EQ(
weight_dims[0],
out_dims[1],
errors::InvalidArgument(
"The second dimension of input(Out@GRAD) must be equal to "
"the third dimension of the Input(Weight)."));
if (dx) {
dx->set_dims(x_dims);
dx->set_dtype(x.dtype());
}
if (dy) {
dy->set_dims(y_dims);
dy->set_dtype(y.dtype());
}
if (dweight) {
dweight->set_dims(weight_dims);
dweight->set_dtype(weight.dtype());
}
if (dbias) {
dbias->set_dims({1, out_dims[1]});
dbias->set_dtype(dout.dtype());
}
}
void BmmGradInferMeta(const MetaTensor& x,
const MetaTensor& y,
const MetaTensor& out_grad,
MetaTensor* x_grad,
MetaTensor* y_grad) {
if (x_grad) {
x_grad->set_dims(x.dims());
x_grad->set_dtype(x.dtype());
}
if (y_grad) {
y_grad->set_dims(y.dims());
y_grad->set_dtype(y.dtype());
}
}
void ChannelShuffleGradInferMeta(const MetaTensor& out_grad,
int groups,
const std::string& data_format,
MetaTensor* x_grad) {
auto do_dims = out_grad.dims();
PADDLE_ENFORCE_EQ(do_dims.size(),
4,
common::errors::InvalidArgument(
"Input should be a 4-D tensor of format [N, C, H, W] "
"or [N, H, W, C], but got %u.",
do_dims.size()));
const auto& dx_dims = do_dims;
x_grad->set_dims(dx_dims);
x_grad->set_dtype(out_grad.dtype());
}
void ComplexGradInferMeta(const MetaTensor& x,
const MetaTensor& y,
const MetaTensor& dout,
MetaTensor* dx,
MetaTensor* dy) {
auto x_dims = x.dims();
if (dx) {
dx->set_dims(x_dims);
dx->set_dtype(x.dtype());
}
auto y_dims = y.dims();
if (dy) {
dy->set_dims(y_dims);
dy->set_dtype(y.dtype());
}
}
void ConvTransposeGradInferMeta(const MetaTensor& x,
const MetaTensor& filter,
const MetaTensor& dout,
const std::vector<int>& strides,
const std::vector<int>& paddings,
const std::vector<int>& output_padding,
const std::vector<int>& output_size,
const std::string& padding_algorithm,
int groups,
const std::vector<int>& dilations,
const std::string& data_format,
MetaTensor* dx,
MetaTensor* dfilter) {
GeneralBinaryGradInferMeta(x, filter, dx, dfilter);
}
void Conv2dTransposeGradInferMeta(const MetaTensor& x,
const MetaTensor& filter,
const MetaTensor& dout,
const std::vector<int>& strides,
const std::vector<int>& paddings,
const std::vector<int>& output_padding,
const IntArray& output_size,
const std::string& padding_algorithm,
int groups,
const std::vector<int>& dilations,
const std::string& data_format,
MetaTensor* dx,
MetaTensor* dfilter) {
GeneralBinaryGradInferMeta(x, filter, dx, dfilter);
}
void Conv2dTransposeDoubleGradInferMeta(const MetaTensor& x,
const MetaTensor& filter,
const MetaTensor& dout,
const MetaTensor& ddx,
const MetaTensor& ddfilter,
const std::vector<int>& strides,
const std::vector<int>& paddings,
const std::vector<int>& output_padding,
const IntArray& output_size,
const std::string& padding_algorithm,
int groups,
const std::vector<int>& dilations,
const std::string& data_format,
MetaTensor* dx,
MetaTensor* dfilter,
MetaTensor* ddout) {
GeneralBinaryGradInferMeta(x, filter, dx, dfilter);
if (ddout) {
ddout->share_meta(dout);
}
}
void CropGradInferMeta(const MetaTensor& out_grad,
const MetaTensor& x,
const IntArray& offsets,
MetaTensor* x_grad) {
auto x_dims = x.dims();
if (x_grad != nullptr) {
x_grad->set_dims(x_dims);
x_grad->set_dtype(x.dtype());
}
}
void CrossEntropyGradInferMeta(const MetaTensor& x,
const MetaTensor& label,
const MetaTensor& out_grad,
bool soft_label,
int ignore_index,
MetaTensor* x_grad,
MetaConfig config) {
const auto& x_dims = x.dims();
const auto& label_dims = label.dims();
const auto& dy_dims = out_grad.dims();
int rank = x_dims.size();
PADDLE_ENFORCE_EQ(dy_dims.size(),
label_dims.size(),
common::errors::InvalidArgument(
"Input(Y@GRAD) and Input(Y) should have the same rank."
"But received: Y@GRAD's rank is [%d], Y's rank is [%d]",
dy_dims.size(),
label_dims.size()));
bool contain_unknown_dim = common::contain_unknown_dim(x_dims) ||
common::contain_unknown_dim(dy_dims);
bool check = config.is_runtime || !contain_unknown_dim;
if (check) {
PADDLE_ENFORCE_EQ(slice_ddim(x_dims, 0, rank - 1),
slice_ddim(dy_dims, 0, rank - 1),
common::errors::InvalidArgument(
"The Input(X) and Input(Y@GRAD) should have the same "
"shape except the last dimension. but received: "
"the shape of Input(X) is [%s], "
"the shape of Input(Y@GRAD) is [%s].",
x_dims,
dy_dims));
}
x_grad->set_dims(x_dims);
x_grad->share_lod(x);
x_grad->set_dtype(x.dtype());
}
void CrossEntropyGrad2InferMeta(const MetaTensor& x_shape,
const MetaTensor& label,
const MetaTensor& match_x,
const MetaTensor& out_grad,
int ignore_index,
MetaTensor* x_grad,
MetaConfig config) {
const auto& x_shape_dims = x_shape.dims();
const auto& x_dims = DDim(x_shape_dims.Get(), x_shape_dims.size() - 1);
const auto& label_dims = label.dims();
const auto& dy_dims = out_grad.dims();
int rank = x_dims.size();
PADDLE_ENFORCE_EQ(dy_dims.size(),
label_dims.size(),
common::errors::InvalidArgument(
"Input(Y@GRAD) and Input(Y) should have the same rank."
"But received: Y@GRAD's rank is [%d], Y's rank is [%d]",
dy_dims.size(),
label_dims.size()));
bool contain_unknown_dim = common::contain_unknown_dim(x_dims) ||
common::contain_unknown_dim(dy_dims);
bool check = config.is_runtime || !contain_unknown_dim;
if (check) {
PADDLE_ENFORCE_EQ(slice_ddim(x_dims, 0, rank - 1),
slice_ddim(dy_dims, 0, rank - 1),
common::errors::InvalidArgument(
"The Input(X) and Input(Y@GRAD) should have the same "
"shape except the last dimension. but received: "
"the shape of Input(X) is [%s], "
"the shape of Input(Y@GRAD) is [%s].",
x_dims,
dy_dims));
}
x_grad->set_dims(x_dims);
x_grad->share_lod(x_shape);
x_grad->set_dtype(x_shape.dtype());
}
void CSoftmaxWithCrossEntropyGradInferMeta(const MetaTensor& softmax,
const MetaTensor& label,
const MetaTensor& loss_grad,
int64_t ignore_index,
int rank,
int nranks,
MetaTensor* logits_grad,
MetaConfig config) {
logits_grad->set_dims(softmax.dims());
}
void CSoftmaxWithMultiLabelCrossEntropyGradInferMeta(
const MetaTensor& softmax,
const MetaTensor& label,
const MetaTensor& smooth_weight,
const MetaTensor& loss_grad,
int64_t ignore_index,
bool sum_multi_label_loss,
int rank,
int nranks,
MetaTensor* logits_grad,
MetaConfig config) {
logits_grad->set_dims(softmax.dims());
}
void FlashAttnGradInferMeta(const MetaTensor& q,
const MetaTensor& k,
const MetaTensor& v,
MetaTensor* dq,
MetaTensor* dk,
MetaTensor* dv) {
if (dq) {
dq->share_meta(q);
}
if (dk && k) {
dk->share_meta(k);
}
if (dv && v) {
dv->share_meta(v);
}
}
void FlashAttnQKVPackedGradInferMeta(const MetaTensor& qkv, MetaTensor* dqkv) {
if (dqkv) {
dqkv->share_meta(qkv);
}
}
void FlashAttnV3GradInferMeta(const MetaTensor& q,
const MetaTensor& k,
const MetaTensor& v,
MetaTensor* dq,
MetaTensor* dk,
MetaTensor* dv) {
if (dq) {
dq->share_meta(q);
}
if (dk) {
dk->share_meta(k);
}
if (dv) {
dv->share_meta(v);
}
}
void FlashAttnV3VarlenGradInferMeta(const MetaTensor& q,
const MetaTensor& k,
const MetaTensor& v,
MetaTensor* dq,
MetaTensor* dk,
MetaTensor* dv) {
if (dq) {
dq->share_meta(q);
}
if (dk) {
dk->share_meta(k);
}
if (dv) {
dv->share_meta(v);
}
}
void Flatten2GradInferMeta(const MetaTensor& x,
const MetaTensor& x_shape,
const MetaTensor& out_grad,
int axis,
MetaTensor* x_grad) {
const auto& xshape_dims = x_shape.dims();
auto x_dims = slice_ddim(xshape_dims, 1, xshape_dims.size());
x_grad->set_dims(x_dims);
x_grad->share_lod(x_shape);
x_grad->set_dtype(out_grad.dtype());
}
void FusedDropoutAddGradInferMeta(const MetaTensor& seed_offset,
const MetaTensor& out_grad,
MetaTensor* x_grad,
MetaTensor* y_grad) {
if (x_grad != nullptr) {
x_grad->share_meta(out_grad);
}
if (y_grad != nullptr) {
y_grad->share_meta(out_grad);
}
}
void CrossEntropyWithSoftmaxGradInferMeta(const MetaTensor& label,
const MetaTensor& softmax,
const MetaTensor& loss_grad,
bool soft_label,
bool use_softmax,
bool numeric_stable_mode,
int ignore_index,
int axis,
MetaTensor* logits_grad,
MetaConfig config) {
auto softmax_dims = softmax.dims();
auto labels_dims = label.dims();
auto softmax_rank = softmax_dims.size();
PADDLE_ENFORCE_GE(axis,
-softmax_rank,
common::errors::InvalidArgument(
"Attr(axis) value should be in range [-R, R-1], "
"R is the rank of Input(Logits)."));
PADDLE_ENFORCE_LT(axis,
softmax_rank,
common::errors::InvalidArgument(
"Attr(axis) value should be in range [-R, R-1], "
"R is the rank of Input(Logits)."));
axis = funcs::CanonicalAxis(axis, softmax_rank);
for (int i = 0; i < softmax_rank; i++) {
if (i != axis) {
if (config.is_runtime || (softmax_dims[i] > 0 && labels_dims[i] > 0)) {
PADDLE_ENFORCE_EQ(
softmax_dims[i],
labels_dims[i],
common::errors::InvalidArgument(
"Input(Logits) and Input(Label) should in same shape in "
"dimensions except axis."));
}
}
}
if (soft_label) {
if (config.is_runtime ||
(softmax_dims[axis] > 0 && labels_dims[axis] > 0)) {
PADDLE_ENFORCE_EQ(softmax_dims[axis],
labels_dims[axis],
common::errors::InvalidArgument(
"If Attr(soft_label) == true, "
"the axis dimension of "
"Input(X) and Input(Label) should be equal."));
}
} else {
if (config.is_runtime || labels_dims[axis] > 0) {
PADDLE_ENFORCE_EQ(
labels_dims[axis],
1UL,
common::errors::InvalidArgument("If Attr(soft_label) == false, "
"the axis dimension of "
"Input(Label) should be 1."));
}
}
logits_grad->set_dims(softmax.dims());
logits_grad->set_dtype(softmax.dtype());
}
void CudnnLSTMGradInferMeta(
const MetaTensor& x,
const MetaTensor& init_h,
const MetaTensor& init_c,
const paddle::optional<std::vector<const MetaTensor*>>& weight_list,
MetaTensor* x_grad,
MetaTensor* init_h_grad,
MetaTensor* init_c_grad,
std::vector<MetaTensor*> weight_list_grad) {
if (x_grad) {
x_grad->share_meta(x);
}
if (init_h_grad) {
init_h_grad->share_meta(init_h);
}
if (init_c_grad) {
init_c_grad->share_meta(init_c);
}
if (!weight_list_grad.empty()) {
UnchangedMultiInferMeta(weight_list.get(), weight_list_grad);
}
}
void LinearV2GradInferMeta(const MetaTensor& input,
const MetaTensor& weight,
const MetaTensor& bias,
const MetaTensor& out_grad,
const bool transpose_weight,
MetaTensor* input_grad,
MetaTensor* weight_grad,
MetaTensor* bias_grad) {
auto input_dims = input.dims();
auto weight_dims = weight.dims();
auto bias_dims = bias.dims();
auto dout_dims = out_grad.dims();
// Assume weight to be [K, N] if not transposed, [N, K] if transposed
const int64_t weight_elewise_dim =
transpose_weight ? weight_dims[0] : weight_dims[1];
auto dout_mat_dims = common::flatten_to_2d(dout_dims, dout_dims.size() - 1);
const int64_t input_ndim = input_dims.size();
auto k_from_dout = input_ndim >= 2 ? dout_dims[input_ndim - 2] : 1;
auto k_from_input = input_ndim >= 2 ? input_dims[input_ndim - 2] : 1;
bool check_k =
(k_from_dout < 0 || k_from_input < 0) || (k_from_dout == k_from_input);
if (check_k) {
PADDLE_ENFORCE_EQ(
dout_mat_dims[1],
weight_elewise_dim,
common::errors::InvalidArgument(
"The last dimension of DOut should be equal with Y's last "
"dimension. But received DOut[-1] = [%d], Y[1] = [%d].",
dout_mat_dims[1],
weight_elewise_dim));
}
for (int32_t i = 0; i + 2 < input_dims.size(); ++i) {
if (dout_dims[i] > 0 && input_dims[i] > 0) {
PADDLE_ENFORCE_EQ(
dout_dims[i],
input_dims[i],
common::errors::InvalidArgument(
"The i dimension of DOut should be equal with i dimension of X."
"But received DOut[%d] = [%d], Y[%d] = [%d].",
i,
dout_dims[i],
i,
input_dims[i]));
}
}
if (input_grad && input) {
input_grad->set_dims(input_dims);
input_grad->set_dtype(input.dtype());
}
if (weight_grad && weight) {
weight_grad->set_dims(weight_dims);
weight_grad->set_dtype(weight.dtype());
}
if (bias_grad && bias) {
bias_grad->set_dims(bias_dims);
bias_grad->set_dtype(bias.dtype());
}
}
void LSTMGradInferMeta(const MetaTensor& input,
const MetaTensor& h0,
const MetaTensor& c0,
const MetaTensor& weight,
const MetaTensor& bias,
MetaTensor* input_grad,
MetaTensor* h0_grad,
MetaTensor* c0_grad,
MetaTensor* weight_grad,
MetaTensor* bias_grad,
MetaConfig config) {
if (input_grad) {
input_grad->share_meta(input);
}
if (h0_grad) {
h0_grad->share_meta(h0);
}
if (c0_grad) {
c0_grad->share_meta(c0);
}
if (weight_grad) {
weight_grad->share_meta(weight);
}
if (bias_grad) {
bias_grad->share_meta(bias);
}
}
void DeformableConvGradInferMeta(const MetaTensor& x,
const MetaTensor& offset,
const MetaTensor& filter,
const MetaTensor& mask,
const MetaTensor& out_grad,
const std::vector<int>& strides,
const std::vector<int>& paddings,
const std::vector<int>& dilations,
int deformable_groups,
int groups,
int im2col_step,
MetaTensor* dx,
MetaTensor* offset_grad,
MetaTensor* filter_grad,
MetaTensor* mask_grad) {
GeneralTernaryGradInferMeta(x, offset, filter, dx, offset_grad, filter_grad);
if (mask) {
UnchangedInferMeta(mask, mask_grad);
}
}
void EigGradInferMeta(const MetaTensor& out_w,
const MetaTensor& out_v,
const MetaTensor& dout_w,
const MetaTensor& dout_v,
MetaTensor* dx) {
auto dims = out_v.dims();
if (dx) {
dx->set_dims(dims);
}
}
void EigvalshGradInferMeta(const MetaTensor& out_v,
const MetaTensor& out_w_grad,
const std::string& uplo,
bool is_test,
MetaTensor* x_grad) {
auto dims = out_v.dims();
if (x_grad != nullptr) {
x_grad->set_dims(dims);
x_grad->set_dtype(out_v.dtype());
}
}
void EmbeddingGradInferMeta(const MetaTensor& x,
const MetaTensor& weight,
MetaTensor* out) {
(void)x;
if (weight) {
out->share_dims(weight);
out->set_dtype(weight.dtype());
}
}
void FFTC2RGradInferMeta(const MetaTensor& x,
const std::vector<int64_t>& axes,
const std::string& normalization,
bool forward,
int64_t last_dim_size,
MetaTensor* out,
MetaConfig config) {
PADDLE_ENFORCE_NOT_NULL(out,
common::errors::InvalidArgument(
"Output of fft_c2r _grad should not be null."));
const DDim x_dim = x.dims();
// only ensure that fft axes' size greater than zero at runtime
// they might be -1 to indicate unknown size ar compile time
if (config.is_runtime) {
for (auto axis : axes) {
PADDLE_ENFORCE_GT(x_dim[axis],
0,
common::errors::InvalidArgument(
"Invalid fft n-point (%d).", x_dim[axis]));
}
}
out->set_layout(x.layout());
out->set_dtype(ToComplexType(x.dtype()));
DDim out_dim = x.dims();
const int last_fft_axis = static_cast<int>(axes.back());
if (last_dim_size > 0) {
out_dim.at(last_fft_axis) = last_dim_size / 2 + 1;
} else if (config.is_runtime) {
const int64_t last_fft_dim_size = x_dim[last_fft_axis];
out_dim.at(last_fft_axis) = last_fft_dim_size / 2 + 1;
} else {
const int64_t last_fft_dim_size = x_dim[last_fft_axis];
out_dim.at(last_fft_axis) =
last_fft_dim_size == -1 ? -1 : last_fft_dim_size / 2 + 1;
}
out->set_dims(out_dim);
}
void FillDiagonalGradInferMeta(const MetaTensor& dout,
float value,
int offset,
bool wrap,
MetaTensor* dx) {
auto x_dims = dout.dims();
if (dx) {
dx->set_dims(x_dims);
dx->set_dtype(dout.dtype());
}
}
void FillDiagonalTensorGradInferMeta(const MetaTensor& out_grad,
int64_t offset,
int dim1,
int dim2,
MetaTensor* x_grad) {
if (x_grad != nullptr) {
x_grad->set_dims(out_grad.dims());
x_grad->set_dtype(out_grad.dtype());
}
}
void GatherNdGradInferMeta(const MetaTensor& x,
const MetaTensor& index,
const MetaTensor& out_grad,
MetaTensor* x_grad) {
const auto& dtype = out_grad.dtype();
x_grad->set_dims(x.dims());
x_grad->share_lod(x);
x_grad->set_dtype(dtype);
}
void GeneralBinaryGradInferMeta(const MetaTensor& x,
const MetaTensor& y,
MetaTensor* dx,
MetaTensor* dy) {
if (dx) {
dx->share_meta(x);
}
if (dy) {
dy->share_meta(y);
}
}
void GeneralTernaryGradInferMeta(const MetaTensor& x,
const MetaTensor& y,
const MetaTensor& z,
MetaTensor* dx,
MetaTensor* dy,
MetaTensor* dz) {
if (dx) {
dx->share_meta(x);
}
if (dy && y) {
dy->share_meta(y);
}
if (dz) {
if (z) {
dz->share_meta(z);
}
}
}
void GeneralQuaternaryGradInferMeta(const MetaTensor& x,
const MetaTensor& y,
const MetaTensor& z,
const MetaTensor& k,
MetaTensor* dx,
MetaTensor* dy,
MetaTensor* dz,
MetaTensor* dk) {
if (dx) {
dx->share_meta(x);
}
if (dy) {
dy->share_meta(y);
}
if (dz) {
dz->share_meta(z);
}
if (dk) {
dk->share_meta(k);
}
}
void GeneralQuinaryGradInferMeta(const MetaTensor& x,
const MetaTensor& y,
const MetaTensor& z,
const MetaTensor& k,
const MetaTensor& l,
MetaTensor* dx,
MetaTensor* dy,
MetaTensor* dz,
MetaTensor* dk,
MetaTensor* dl) {
if (dx) {
dx->share_meta(x);
}
if (dy) {
dy->share_meta(y);
}
if (dz) {
dz->share_meta(z);
}
if (dk) {
dk->share_meta(k);
}
if (dl) {
dl->share_meta(l);
}
}
void GruGradInferMeta(const MetaTensor& input,
const MetaTensor& h0,
const MetaTensor& weight,
const MetaTensor& bias,
MetaTensor* input_grad,
MetaTensor* h0_grad,
MetaTensor* weight_grad,
MetaTensor* bias_grad,
MetaConfig config) {
const auto& input_dims = input.dims();
const auto& weight_dims = weight.dims();
int64_t input_size = input_dims[1];
int64_t frame_size = weight_dims[0];
int64_t weight_height = weight_dims[0];
int64_t weight_width = weight_dims[1];
PADDLE_ENFORCE_EQ(
input_size,
frame_size * 3,
common::errors::InvalidArgument(
"The second dimension of Input(Input) must be 3 times of "
"frame_size in GRUOp, but received %d (Input) vs %d (frame_size).",
input_size,
frame_size));
PADDLE_ENFORCE_EQ(
weight_height,
frame_size,
common::errors::InvalidArgument(
"The shape of Input(Weight) matrix must be [frame_size, frame_size "
"* 3], but received [%d, %d] (Weight) vs [%d, %d] (frame_size).",
weight_height,
weight_width,
frame_size,
frame_size * 3));
PADDLE_ENFORCE_EQ(
weight_width,
frame_size * 3,
common::errors::InvalidArgument(
"The shape of Input(Weight) matrix must be [frame_size, frame_size "
"* 3], but received [%d, %d] (Weight) vs [%d, %d] (frame_size).",
weight_height,
weight_width,
frame_size,
frame_size * 3));
if (h0.initialized()) {
const auto& h0_dims = h0.dims();
PADDLE_ENFORCE_EQ(
h0_dims[1],
frame_size,
common::errors::InvalidArgument(
"The width of Input(H0) must be equal to frame_size, but "
"received %d (width of H0) vs %d (frame_size).",
h0_dims[1],
frame_size));
if (h0_grad != nullptr) {
h0_grad->set_dims(h0_dims);
h0_grad->set_dtype(h0.dtype());
}
}
if (bias.initialized()) {
const auto& bias_dims = bias.dims();
int64_t bias_height = bias_dims[0];
int64_t bias_width = bias_dims[1];
PADDLE_ENFORCE_EQ(
bias_height,
1,
common::errors::InvalidArgument(
"The shape of Bias must be [1, frame_size * 3], but received "
"[%d, %d] (Bias) vs [1, %d] (frame_size * 3).",
bias_height,
bias_width,
frame_size * 3));
PADDLE_ENFORCE_EQ(
bias_width,
frame_size * 3,
common::errors::InvalidArgument(
"The shape of Bias must be [1, frame_size * 3], but received "
"[%d, %d] (Bias) vs [1, %d] (frame_size * 3).",
bias_height,
bias_width,
frame_size * 3));
if (bias_grad != nullptr) {
bias_grad->set_dims(bias_dims);
bias_grad->set_dtype(bias.dtype());
}
}
if (input_grad != nullptr) {
input_grad->set_dims(input_dims);
input_grad->set_dtype(input.dtype());
}
if (weight_grad != nullptr) {
weight_grad->set_dims(weight_dims);
weight_grad->set_dtype(weight.dtype());
}
}
void GruUnitGradInferMeta(const MetaTensor& input,
const MetaTensor& hidden_prev,
const MetaTensor& weight,
const MetaTensor& bias,
MetaTensor* input_grad,
MetaTensor* hidden_prev_grad,
MetaTensor* weight_grad,
MetaTensor* bias_grad,
MetaConfig config) {
const auto& input_dims = input.dims();
const auto& hidden_prev_dims = hidden_prev.dims();
const auto& weight_dims = weight.dims();
int64_t input_size = input_dims[1];
int64_t frame_size = hidden_prev_dims[1];
int64_t weight_height = weight_dims[0];
int64_t weight_width = weight_dims[1];
if (config.is_runtime || input_size >= 0) {
PADDLE_ENFORCE_EQ(
input_size,
frame_size * 3,
common::errors::InvalidArgument(
"The second dimension of Input(Input) must be 3 "
"times of frame_size in GRUUnitGradOp, but received %d "
"(Input) vs %d (frame_size).",
input_size,
frame_size));
}
PADDLE_ENFORCE_EQ(
weight_height,
frame_size,
common::errors::InvalidArgument(
"The shape of Input(Weight) matrix must be [frame_size, frame_size "
"* 3] in GRUUnitGradOp, but received [%d, %d] (Weight) vs [%d, %d] "
"(frame_size).",
weight_height,
weight_width,
frame_size,
frame_size * 3));
PADDLE_ENFORCE_EQ(
weight_width,
frame_size * 3,
common::errors::InvalidArgument(
"The shape of Input(Weight) matrix must be [frame_size, frame_size "
"* 3] in GRUUnitGradOp, but received [%d, %d] (Weight) vs [%d, %d] "
"(frame_size).",
weight_height,
weight_width,
frame_size,
frame_size * 3));
if (bias.initialized()) {
const auto& bias_dims = bias.dims();
int64_t bias_height = bias_dims[0];
int64_t bias_width = bias_dims[1];
PADDLE_ENFORCE_EQ(
bias_height,
1,
common::errors::InvalidArgument(
"The shape of Bias must be [1, frame_size * 3], but received "
"[%d, %d] (Bias) vs [1, %d] (frame_size * 3).",
bias_height,
bias_width,
frame_size * 3));
PADDLE_ENFORCE_EQ(
bias_width,
frame_size * 3,
common::errors::InvalidArgument(
"The shape of Bias must be [1, frame_size * 3], but received "
"[%d, %d] (Bias) vs [1, %d] (frame_size * 3).",
bias_height,
bias_width,
frame_size * 3));
if (bias_grad != nullptr) {
bias_grad->set_dims(bias_dims);
bias_grad->set_dtype(bias.dtype());
}
}
if (input_grad != nullptr) {
input_grad->set_dims(input_dims);
input_grad->set_dtype(input.dtype());
}
if (hidden_prev_grad != nullptr) {
hidden_prev_grad->set_dims(hidden_prev_dims);
hidden_prev_grad->set_dtype(hidden_prev.dtype());
}
if (weight_grad != nullptr) {
weight_grad->set_dims(weight_dims);
weight_grad->set_dtype(weight.dtype());
}
}
void GeneralUnaryGradInferMeta(const MetaTensor& x, MetaTensor* dx) {
if (dx) {
dx->share_meta(x);
}
}
void GumbelSoftmaxGradInferMeta(const MetaTensor& out,
const MetaTensor& dout,
int axis,
MetaTensor* dx) {
PADDLE_ENFORCE_EQ(
out.dims(),
dout.dims(),
errors::InvalidArgument(
"Input(Out) and its gradients should have the same shape."));
dx->share_meta(dout);
}
void InstanceNormGradInferMeta(const MetaTensor& x,
const MetaTensor& scale,
const MetaTensor& bias,
const MetaTensor& saved_mean,
const MetaTensor& saved_variance,
const MetaTensor& y_grad,
float epsilon,
MetaTensor* x_grad,
MetaTensor* scale_grad,
MetaTensor* bias_grad) {
PADDLE_ENFORCE_NE(
x_grad,
nullptr,
common::errors::InvalidArgument(
"The X@GRAD in InstanceNormGradInferMeta can't be nullptr."));
const auto x_dims = x.dims();
const int64_t C = x_dims[1];
x_grad->set_dims(x_dims);
x_grad->set_dtype(x.dtype());
x_grad->set_layout(x.layout());
if (scale_grad) {
if (C == 0) {
scale_grad->set_dims({scale.dims()[0]});
} else {
scale_grad->set_dims({C});
}
}
if (bias_grad) {
if (C == 0) {
bias_grad->set_dims({bias.dims()[0]});
} else {
bias_grad->set_dims({C});
}
}
}
void InstanceNormDoubleGradInferMeta(const MetaTensor& x,
const MetaTensor& scale,
const MetaTensor& saved_mean,
const MetaTensor& saved_variance,
const MetaTensor& dy,
const MetaTensor& ddx,
const MetaTensor& ddscale,
const MetaTensor& ddbias,
float epsilon,
MetaTensor* dx,
MetaTensor* dscale,
MetaTensor* ddy) {
PADDLE_ENFORCE_NE(
dx,
nullptr,
common::errors::InvalidArgument(
"The DX in InstanceNormDoubleGradInferMeta can't be nullptr."));
const auto x_dims = x.dims();
const int64_t C = x_dims[1];
dx->set_dims(x_dims);
dx->set_dtype(x.dtype());
dx->set_layout(x.layout());
if (dscale) {
dscale->set_dims({C});
}
if (ddy) {
ddy->share_dims(x);
}
}
void InverseGradInferMeta(const MetaTensor& out,
const MetaTensor& dout,
MetaTensor* dx) {
if (dx) {
dx->set_dims(dout.dims());
dx->set_dtype(out.dtype());
}
}
void KernelWithXShapeInferMeta(const MetaTensor& xshape,
const MetaTensor& out,
MetaTensor* dx) {
auto xshape_dims = xshape.dims();
auto x_dims = slice_ddim(xshape_dims, 1, xshape_dims.size());
dx->set_dims(x_dims);
dx->set_dtype(out.dtype());
dx->share_lod(xshape);
}
void GradSameWithXInferMeta(const MetaTensor& x,
const MetaTensor& out,
MetaTensor* dx) {
dx->set_dims(x.dims());
dx->set_dtype(out.dtype());
dx->share_lod(x);
}
void LodResetGradInferMeta(const MetaTensor& x,
const MetaTensor& out_grad,
const std::vector<int>& target_lod,
bool append,
MetaTensor* x_grad,
MetaConfig config) {
if (x_grad != nullptr) {
x_grad->set_dims(x.dims());
x_grad->share_lod(x);
x_grad->set_dtype(x.dtype());
}
}
void LUGradInferMeta(const MetaTensor& x,
const MetaTensor& out,
const MetaTensor& pivots,
const MetaTensor& out_grad,
bool pivot,
MetaTensor* x_grad) {
auto x_dims = x.dims();
if (x_grad) {
x_grad->set_dims(x_dims);
x_grad->set_dtype(x.dtype());
}
}
void LUUnpackGradInferMeta(const MetaTensor& x,
const MetaTensor& pivots,
const MetaTensor& l,
const MetaTensor& u,
const MetaTensor& pmat,
const MetaTensor& l_grad,
const MetaTensor& u_grad,
bool unpack_ludata,
bool unpack_pivots,
MetaTensor* x_grad) {
auto x_dims = x.dims();
if (x_grad) {
x_grad->set_dims(x_dims);
x_grad->set_dtype(x.dtype());
}
}
void MarginCrossEntropyGradInferMeta(const MetaTensor& logits,
const MetaTensor& label,
const MetaTensor& softmax,
const MetaTensor& loss_grad,
bool return_softmax,
int ring_id,
int rank,
int nranks,
float margin1,
float margin2,
float margin3,
float scale,
MetaTensor* logits_grad) {
PADDLE_ENFORCE_NE(
logits_grad,
nullptr,
common::errors::InvalidArgument(
"The Logits@GRAD in MarginCrossEntropy can't be nullptr."));
auto softmax_dims = softmax.dims();
logits_grad->set_dims(softmax_dims);
logits_grad->set_dtype(softmax.dtype());
}
void MatchMatrixTensorGradInferMeta(const MetaTensor& x,
const MetaTensor& y,
const MetaTensor& w,
const MetaTensor& tmp,
const MetaTensor& out_grad,
int dim_t,
MetaTensor* x_grad,
MetaTensor* y_grad,
MetaTensor* w_grad) {
if (x_grad != nullptr) {
x_grad->set_dims(x.dims());
x_grad->share_lod(x);
x_grad->set_dtype(x.dtype());
}
if (y_grad != nullptr) {
y_grad->set_dims(y.dims());
y_grad->share_lod(y);
y_grad->set_dtype(y.dtype());
}
if (w_grad != nullptr) {
w_grad->set_dims(w.dims());
w_grad->set_dtype(w.dtype());
}
}
void MaxPoolWithIndexGradInferMeta(const MetaTensor& x,
const MetaTensor& mask,
const MetaTensor& dout,
const std::vector<int>& kernel_size,
const std::vector<int>& strides,
const std::vector<int>& paddings,
const std::vector<int>& dilations,
bool global_pooling,
bool adaptive,
bool ceil_mode,
MetaTensor* dx) {
dx->share_meta(x);
}
void MedianGradInferMeta(const MetaTensor& x,
const MetaTensor& median_data,
const MetaTensor& median_index,
const MetaTensor& out_grad,
const IntArray& axes,
bool keep_dim,
const std::string& mode,
MetaTensor* x_grad) {
auto x_dims = x.dims();
x_grad->set_dims(x_dims);
x_grad->set_dtype(x.dtype());
}
void MemoryEfficientAttentionGradInferMeta(const MetaTensor& query,
const MetaTensor& key,
const MetaTensor& value,
const MetaTensor& bias,
const MetaTensor& cu_seqlens_q,
const MetaTensor& cu_seqlens_k,
const MetaTensor& output,
const MetaTensor& logsumexp,
const MetaTensor& seed_and_offset,
const MetaTensor& output_grad,
const Scalar& max_seqlen_q,
const Scalar& max_seqlen_k,
const bool causal,
const double dropout_p,
const float scale,
MetaTensor* query_grad,
MetaTensor* key_grad,
MetaTensor* value_grad,
MetaTensor* bias_grad) {
PADDLE_ENFORCE_EQ(
output_grad.dims().size(),
4,
common::errors::InvalidArgument("Key should be a 4-D tensor. "
"But received Key dimension(%s)",
output_grad.dims().size()));
PADDLE_ENFORCE_EQ(
output.dims().size(),
4,
common::errors::InvalidArgument("Key should be a 4-D tensor. "
"But received Key dimension(%s)",
output_grad.dims().size()));
const int64_t query_batch_size = query.dims()[0];
const int64_t query_seq_length = query.dims()[1];
const int64_t query_num_head = query.dims()[2];
const int64_t query_head_size = query.dims()[3];
const int64_t key_batch_size = key.dims()[0];
const int64_t key_seq_length = key.dims()[1];
const int64_t key_num_head = key.dims()[2];
const int64_t key_head_size = key.dims()[3];
const int64_t value_batch_size = value.dims()[0];
const int64_t value_seq_length = value.dims()[1];
const int64_t value_num_head = value.dims()[2];
const int64_t value_head_size = value.dims()[3];
if (query_grad) {
std::vector<int64_t> query_grad_dims;
query_grad_dims = {
query_batch_size, query_seq_length, query_num_head, query_head_size};
query_grad->set_dims(make_ddim(query_grad_dims));
query_grad->share_lod(query);
query_grad->set_dtype(query.dtype());
query_grad->set_layout(query.layout());
}
if (key_grad) {
std::vector<int64_t> key_grad_dims;
key_grad_dims = {
key_batch_size, key_seq_length, key_num_head, key_head_size};
key_grad->set_dims(make_ddim(key_grad_dims));
key_grad->share_lod(key);
key_grad->set_dtype(key.dtype());
key_grad->set_layout(key.layout());
}
if (value_grad) {
std::vector<int64_t> value_grad_dims;
value_grad_dims = {
value_batch_size, value_seq_length, value_num_head, value_head_size};
value_grad->set_dims(make_ddim(value_grad_dims));
value_grad->share_lod(value);
value_grad->set_dtype(value.dtype());
value_grad->set_layout(value.layout());
}
if (bias && bias_grad) {
const int64_t bias_batch_size = bias.dims()[0];
const int64_t bias_seq_length = bias.dims()[1];
const int64_t bias_num_head = bias.dims()[2];
const int64_t bias_head_size = bias.dims()[3];
std::vector<int64_t> bias_grad_dims(
{bias_batch_size, bias_seq_length, bias_num_head, bias_head_size});
bias_grad->set_dims(make_ddim(bias_grad_dims));
bias_grad->share_lod(bias);
bias_grad->set_dtype(bias.dtype());
bias_grad->set_layout(bias.layout());
} else if (bias_grad) {
std::vector<int64_t> bias_grad_dims;
bias_grad->set_dims(make_ddim(bias_grad_dims));
}
}
void MeshgridGradInferMeta(const std::vector<const MetaTensor*>& inputs,
const std::vector<const MetaTensor*>& outputs_grad,
std::vector<MetaTensor*> inputs_grad) {
PADDLE_ENFORCE_GT(outputs_grad.size(),
1,
errors::InvalidArgument(
"Number of Inputs(Out@GRAD) should be larger than 1."
"But received Inputs(Out@GRAD)' size = %d .",
outputs_grad.size()));
for (size_t i = 0; i < inputs.size(); i++) {
inputs_grad[i]->share_meta(*inputs[i]);
}
}
void MoeCombineGradInferMeta(const MetaTensor& x,
const MetaTensor& combine_weights,
const MetaTensor& scatter_index,
const MetaTensor& y,
MetaTensor* grad_x,
MetaTensor* grad_combine_weights_helper) {
auto x_dim = x.dims();
auto combine_weights_shape = combine_weights.dims();
PADDLE_ENFORCE_EQ(
x_dim.size(),
2,
errors::InvalidArgument("The input X should have 2 dimensions. "
"But received X's dimension = %d",
x_dim.size()));
PADDLE_ENFORCE_EQ(
(scatter_index.dtype() == DataType::INT32),
true,
errors::InvalidArgument("The input scatter_index type should be int32. "
"But received scatter_index type = %s",
scatter_index.dtype()));
grad_x->set_dims(make_ddim({x_dim[0], x_dim[1]}));
grad_x->set_dtype(x.dtype());
grad_combine_weights_helper->set_dims(make_ddim(
{combine_weights_shape[0], combine_weights_shape[1], x_dim[1]}));
grad_combine_weights_helper->set_dtype(x.dtype());
}
void MoeCombineAutoGradInferMeta(const MetaTensor& x,
const MetaTensor& combine_weights,
const MetaTensor& scatter_index,
const MetaTensor& y,
MetaTensor* grad_x,
MetaTensor* grad_combine_weights_helper,
MetaTensor* grad_scatter_index) {
auto x_dim = x.dims();
auto combine_weights_shape = combine_weights.dims();
auto scatter_index_dim = scatter_index.dims();
PADDLE_ENFORCE_EQ(
x_dim.size(),
2,
errors::InvalidArgument("The input X should have 2 dimensions."
"But received X's dimension = %d",
x_dim.size()));
PADDLE_ENFORCE_EQ(
(scatter_index.dtype() == DataType::INT32),
true,
errors::InvalidArgument("The input scatter_index type should be int32."
"But received scatter_index type = %s",
scatter_index.dtype()));
grad_x->set_dims(make_ddim({x_dim[0], x_dim[1]}));
grad_x->set_dtype(x.dtype());
grad_combine_weights_helper->set_dims(
make_ddim({combine_weights_shape[0], combine_weights_shape[1]}));
grad_combine_weights_helper->set_dtype(x.dtype());
PADDLE_ENFORCE_NE(
grad_scatter_index,
nullptr,
common::errors::InvalidArgument(
"The scatter_index need grad in auto parallel version moe_combine, "
"set scatter_index.stop_gradient = False."));
grad_scatter_index->set_dims(scatter_index_dim);
grad_scatter_index->set_dtype(DataType::INT32);
}
void MoeGateDispatchPartialNoSoftmaxTopkGradInferMeta(
const MetaTensor& combine_weights_out,
const MetaTensor& scatter_index,
const MetaTensor& scatter_index_rev,
const MetaTensor& expert_offset,
const MetaTensor& expert_offset_local,
const MetaTensor& y_grad,
const MetaTensor& combine_weights_out_grad,
int64_t k,
int64_t capacity,
bool use_pad,
int64_t expert_start_index,
int64_t expert_end_index,
MetaTensor* x_grad,
MetaTensor* combine_weights_grad) {
int64_t num_experts = expert_offset.dims()[0];
int64_t hidden_size = y_grad.dims()[1];
int64_t num_rows = scatter_index.dims()[1];
PADDLE_ENFORCE_GT(num_experts,
0,
common::errors::InvalidArgument(
"Input num_experts should be greater than 0"));
PADDLE_ENFORCE_EQ((expert_offset.dtype() == DataType::INT64),
true,
common::errors::InvalidArgument(
"Input expert_offset type should be int64"));
if (use_pad) {
PADDLE_ENFORCE_GE(num_experts,
y_grad.dims()[0] / capacity,
common::errors::InvalidArgument(
"Number of experts should be greater than or equal "
"to y_grad.dims()[0]/capacity"));
} else {
PADDLE_ENFORCE_GT(y_grad.dims()[0],
0,
common::errors::InvalidArgument(
"Input y_grad.dims()[0] should be greater than 0"));
}
combine_weights_grad->set_dims(combine_weights_out_grad.dims());
combine_weights_grad->set_dtype(DataType::FLOAT32);
x_grad->set_dims({num_rows, hidden_size});
x_grad->set_dtype(y_grad.dtype());
}
void MoeGateDispatchPermuteGradInferMeta(const MetaTensor& combine_weights,
const MetaTensor& scatter_index,
const MetaTensor& expert_id,
const MetaTensor& y_grad,
const MetaTensor& combine_weights_grad,
int64_t k,
int64_t capacity,
int64_t world_size,
MetaTensor* x_grad,
MetaTensor* gate_logits_grad) {
auto y_grad_dims = y_grad.dims();
PADDLE_ENFORCE_EQ(
y_grad_dims[1],
world_size,
common::errors::InvalidArgument(
"The second dimension of y_grad should be equal to world_size, but "
"received y_grad_dims[1] = %d, world_size = %d",
y_grad_dims[1],
world_size));
int64_t num_local_experts = y_grad_dims[0];
int64_t num_experts = world_size * num_local_experts;
int64_t hidden_size = y_grad_dims[y_grad_dims.size() - 1];
int64_t num_rows = scatter_index.dims()[1];
x_grad->set_dims({num_rows, hidden_size});
x_grad->set_dtype(y_grad.dtype());
gate_logits_grad->set_dims({num_rows, num_experts});
gate_logits_grad->set_dtype(DataType::FLOAT32);
}
void MultiDotGradInferMeta(const std::vector<const MetaTensor*>& x,
const MetaTensor& out_grad,
std::vector<MetaTensor*> x_grad) {
PADDLE_ENFORCE_EQ(
x.size(),
x_grad.size(),
errors::InvalidArgument(
"Number of Inputs(X) should be equal with Outputs(X@GRAD)."
"But received Inputs(X)' size = %d , Outputs(X@GRAD)' size = %d.",
x.size(),
x_grad.size()));
for (size_t i = 0; i < x.size(); i++) {
if (x_grad[i] != nullptr) {
x_grad[i]->set_dims(x[i]->dims());
x_grad[i]->share_lod(*x[i]);
}
}
}
void MultiplexGradInferMeta(const MetaTensor& ids,
const MetaTensor& out_grad,
std::vector<MetaTensor*> ins_grad) {
PADDLE_ENFORCE_NE(
ins_grad.empty(),
true,
errors::InvalidArgument("Output(X@GRAD) should not be null."));
auto dout_dim = out_grad.dims();
for (auto in_grad : ins_grad) {
if (in_grad != nullptr) {
in_grad->set_dims(dout_dim);
}
}
}
void NanmedianGradInferMeta(const MetaTensor& x,
const MetaTensor& median_data,
const MetaTensor& median_index,
const MetaTensor& out_grad,
const IntArray& axes,
bool keep_dim,
const std::string& mode,
MetaTensor* x_grad) {
auto x_dims = x.dims();
x_grad->set_dims(x_dims);
x_grad->set_dtype(x.dtype());
}
void PartialConcatGradInferMeta(const std::vector<const MetaTensor*>& xs,
std::vector<MetaTensor*> x_grads) {
auto input_num = xs.size();
for (size_t i = 0; i < input_num; i++) {
auto x_dims = xs[i]->dims();
x_grads[i]->set_dims(x_dims);
x_grads[i]->set_dtype(xs[i]->dtype());
}
}
void NceGradInferMeta(const MetaTensor& input,
const MetaTensor& bias,
const MetaTensor& weight,
MetaTensor* input_grad,
MetaTensor* bias_grad,
MetaTensor* weight_grad
) {
auto x_dims = input.dims();
if (input_grad != nullptr) {
input_grad->set_dims(x_dims);
input_grad->set_dtype(input.dtype());
}
auto w_dims = weight.dims();
if (weight_grad) {
weight_grad->set_dims(w_dims);
weight_grad->set_dtype(weight.dtype());
}
auto bias_dims = bias.dims();
if (bias_grad) {
bias_grad->set_dims(bias_dims);
bias_grad->set_dtype(bias.dtype());
}
}
void PartialSumGradInferMeta(const std::vector<const MetaTensor*>& xs,
std::vector<MetaTensor*> x_grads) {
auto input_num = xs.size();
for (size_t i = 0; i < input_num; i++) {
auto x_dims = xs[i]->dims();
x_grads[i]->set_dims(x_dims);
x_grads[i]->set_dtype(xs[i]->dtype());
}
}
void NllLossGradInferMeta(const MetaTensor& x,
const MetaTensor& label,
const MetaTensor& weight,
const MetaTensor& total_weight,
const MetaTensor& out_grad,
int64_t ignore_index,
const std::string& reduction,
MetaTensor* dx,
MetaConfig config) {
const auto& x_dims = x.dims();
const auto& label_dims = label.dims();
const auto& dout_dims = out_grad.dims();
bool contain_unknown_dim = common::contain_unknown_dim(x_dims) ||
common::contain_unknown_dim(dout_dims);
bool check = config.is_runtime || !contain_unknown_dim;
if (check) {
auto batch_size = x_dims[0];
if (x_dims.size() == 2) {
if (reduction == "none") {
PADDLE_ENFORCE_EQ(dout_dims.size(),
1,
common::errors::InvalidArgument(
"The dimensions of Input(Out@GRAD) must be 1"));
PADDLE_ENFORCE_EQ(
dout_dims[0],
batch_size,
common::errors::InvalidArgument(
"The unreduced size ofInput(Out@GRAD) must be the "
"same as batch_size."));
} else {
PADDLE_ENFORCE_EQ(dout_dims.size(),
0,
common::errors::InvalidArgument(
"The dimensions of Input(Out@GRAD) must be 0"));
}
} else if (x_dims.size() == 4) {
if (reduction == "none") {
PADDLE_ENFORCE_EQ(
dout_dims.size(),
3,
common::errors::InvalidArgument(
"The dimensions of Input(Out@GRAD) must be 3,But got [%s].",
dout_dims.size()));
PADDLE_ENFORCE_EQ(dout_dims[0] == label_dims[0] &&
dout_dims[1] == label_dims[1] &&
dout_dims[2] == label_dims[2],
true,
common::errors::InvalidArgument(
"The dimensions of Input(Out@GRAD) must be match "
"to Input(Label) dimensions."));
} else {
PADDLE_ENFORCE_EQ(dout_dims.size(),
0,
common::errors::InvalidArgument(
"The dimensions of Input(Out@GRAD) must be 0"));
}
}
}
if (dx) {
dx->set_dims(x_dims);
dx->set_dtype(x.dtype());
}
}
void OverlapAddGradInferMeta(const MetaTensor& x,
const MetaTensor& out_grad,
int hop_length,
int axis,
MetaTensor* x_grad) {
const auto x_dims = x.dims();
if (x_grad != nullptr) {
x_grad->set_dims(x_dims);
x_grad->set_dtype(x.dtype());
}
}
inline int64_t HandleDynamicDim(int64_t maybe_dynamic_dim,
int64_t static_result) {
return maybe_dynamic_dim == -1 ? -1 : static_result;
}
void PixelUnshuffleGradInferMeta(const MetaTensor& out_grad,
int downscale_factor,
const std::string& data_format,
MetaTensor* x_grad) {
auto do_dims = out_grad.dims();
PADDLE_ENFORCE_EQ(do_dims.size(),
4,
common::errors::InvalidArgument(
"Input should be a 4-D tensor of format [N, C, H, W] "
"or [N, H, W, C], but got %u.",
do_dims.size()));
const bool channel_last = (data_format == "NHWC");
auto dx_dims = do_dims;
dx_dims[0] = do_dims[0];
if (!channel_last) {
dx_dims[1] = HandleDynamicDim(
do_dims[1], do_dims[1] / (downscale_factor * downscale_factor));
dx_dims[2] = HandleDynamicDim(do_dims[2], do_dims[2] * downscale_factor);
dx_dims[3] = HandleDynamicDim(do_dims[3], do_dims[3] * downscale_factor);
} else {
dx_dims[1] = HandleDynamicDim(do_dims[1], do_dims[1] * downscale_factor);
dx_dims[2] = HandleDynamicDim(do_dims[2], do_dims[2] * downscale_factor);
dx_dims[3] = HandleDynamicDim(
do_dims[3], do_dims[3] / (downscale_factor * downscale_factor));
}
x_grad->set_dims(dx_dims);
x_grad->set_dtype(out_grad.dtype());
}
void PreluGradInferMeta(const MetaTensor& x,
const MetaTensor& y,
MetaTensor* dx,
MetaTensor* dy) {
if (dx) {
dx->share_dims(x);
}
if (dy) {
dy->share_dims(y);
}
}
void PsroiPoolGradInferMeta(const MetaTensor& x,
const MetaTensor& rois,
const MetaTensor& rois_num,
const MetaTensor& dout,
int pooled_height,
int pooled_width,
int output_channels,
float spatial_scale,
MetaTensor* dx) {
dx->share_meta(x);
}
void RankAttentionGradInferMeta(const MetaTensor& x,
const MetaTensor& rank_offset,
const MetaTensor& rank_param,
const MetaTensor& input_help,
const MetaTensor& ins_rank,
const MetaTensor& out_grad,
int max_rank,
int max_size,
MetaTensor* rank_param_grad) {
rank_param_grad->set_dims(rank_param.dims());
rank_param_grad->set_dtype(rank_param.dtype());
}
void RealAndImagGradInferMeta(const MetaTensor& out_grad, MetaTensor* dx) {
dx->set_dims(out_grad.dims());
dx->set_dtype(dtype::ToComplex(out_grad.dtype()));
dx->set_layout(out_grad.layout());
}
void ReshapeDoubleGradInferMeta(const MetaTensor& out_grad,
const MetaTensor& x_grad_grad,
MetaTensor* out_grad_grad) {
if (out_grad_grad != nullptr) {
out_grad_grad->share_dims(out_grad);
}
}
void FusedRmsNormQuantGradInferMeta(const MetaTensor& x,
const MetaTensor& norm_weight,
const MetaTensor& norm_bias,
MetaTensor* x_grad,
MetaTensor* norm_weight_grad,
MetaTensor* norm_bias_grad) {
if (x_grad) {
x_grad->share_meta(x);
}
if (norm_weight && norm_weight_grad) {
norm_weight_grad->share_meta(norm_weight);
}
if (norm_bias && norm_bias_grad) {
norm_bias_grad->share_meta(norm_bias);
}
}
PADDLE_API void RMSNormGradInferMeta(
const MetaTensor& x,
const MetaTensor& scale,
const MetaTensor& invvar,
const MetaTensor& y_grad,
const std::vector<int64_t>& normalized_shape,
double epsilon,
MetaTensor* x_grad,
MetaTensor* scale_grad) {
if (x_grad && x) {
x_grad->share_meta(x);
}
if (scale_grad && scale) {
scale_grad->share_meta(scale);
}
}
void RnnGradInferMeta(const MetaTensor& x,
const std::vector<const MetaTensor*>& pre_state,
const std::vector<const MetaTensor*>& weight_list,
MetaTensor* x_grad,
std::vector<MetaTensor*> pre_state_grad,
std::vector<MetaTensor*> weight_grad_list) {
PADDLE_ENFORCE_GT(
pre_state.size(),
0UL,
common::errors::InvalidArgument(
"The input pre_state in RnnGradInferMeta can't be empty."));
PADDLE_ENFORCE_GT(
weight_grad_list.size(),
0UL,
common::errors::InvalidArgument(
"The input weight_grad_list in RnnGradInferMeta can't be empty."));
if (x_grad) {
UnchangedInferMeta(x, x_grad);
}
if (!pre_state_grad.empty()) {
UnchangedMultiInferMeta(pre_state, pre_state_grad);
}
if (!weight_grad_list.empty()) {
UnchangedMultiInferMeta(weight_list, weight_grad_list);
}
}
void RowConvGradInferMeta(const MetaTensor& out_grad,
const MetaTensor& filter,
MetaTensor* x_grad,
MetaTensor* filter_grad) {
if (x_grad != nullptr) {
x_grad->set_dims(out_grad.dims());
}
if (filter_grad != nullptr) {
filter_grad->set_dims(filter.dims());
}
}
void ScatterGradInferMeta(const MetaTensor& index,
const MetaTensor& updates,
const MetaTensor& out_grad,
bool overwrite,
MetaTensor* x_grad,
MetaTensor* updates_grad) {
const auto& dtype = out_grad.dtype();
if (updates_grad) {
updates_grad->set_dims(updates.dims());
updates_grad->set_dtype(dtype);
}
if (x_grad) {
x_grad->set_dims(out_grad.dims());
x_grad->set_dtype(dtype);
}
}
void ScatterNdAddGradInferMeta(const MetaTensor& index,
const MetaTensor& updates,
const MetaTensor& out_grad,
MetaTensor* x_grad,
MetaTensor* updates_grad) {
const auto& dtype = out_grad.dtype();
if (updates_grad) {
updates_grad->set_dims(updates.dims());
updates_grad->set_dtype(dtype);
}
if (x_grad) {
x_grad->set_dims(out_grad.dims());
x_grad->set_dtype(dtype);
}
}
void SequenceConvGradInferMeta(const MetaTensor& x,
const MetaTensor& padding_data,
const MetaTensor& filter,
const MetaTensor& out_grad,
int context_length,
bool padding_trainable,
int context_start,
int context_stride,
MetaTensor* x_grad,
MetaTensor* padding_data_grad,
MetaTensor* filter_grad) {
if (padding_trainable && padding_data_grad != nullptr) {
padding_data_grad->set_dims(padding_data.dims());
padding_data_grad->set_dtype(padding_data.dtype());
}
if (x_grad != nullptr) {
x_grad->set_dims(x.dims());
x_grad->share_lod(x);
x_grad->set_dtype(x.dtype());
}
if (filter_grad != nullptr) {
filter_grad->set_dims(filter.dims());
filter_grad->set_dtype(filter.dtype());
}
}
void ShuffleBatchGradInferMeta(const MetaTensor& shuffle_idx,
const MetaTensor& out_grad,
int startup_seed,
MetaTensor* x_grad) {
x_grad->share_dims(out_grad);
x_grad->share_lod(out_grad);
x_grad->set_dtype(out_grad.dtype());
}
void SpectralNormGradInferMeta(const MetaTensor& weight,
const MetaTensor& u,
const MetaTensor& v,
const MetaTensor& out_grad,
int dim,
int power_iters,
float eps,
MetaTensor* weight_grad) {
auto dim_x = weight.dims();
if (weight_grad) {
weight_grad->set_dims(dim_x);
weight_grad->set_dtype(out_grad.dtype());
}
}
void StackGradInferMeta(const MetaTensor& out_grad,
int axis,
std::vector<MetaTensor*> x_grad) {
auto dy_dim = out_grad.dims();
int rank = dy_dim.size();
PADDLE_ENFORCE_GE(
axis,
-rank,
common::errors::InvalidArgument(
"Attr(axis) must be inside [-rank, rank), where rank = %d, "
"but received axis is:%d.",
rank,
axis));
PADDLE_ENFORCE_LT(
axis,
rank,
common::errors::InvalidArgument(
"Attr(axis) must be inside [-rank, rank), where rank = %d, "
"but received axis is:%d.",
rank,
axis));
if (axis < 0) axis += rank;
PADDLE_ENFORCE_LE(
x_grad.size(),
static_cast<size_t>(dy_dim[axis]),
common::errors::InvalidArgument(
"Number of Outputs(X@GRAD) should be less than or equal to dy dim "
"at axis, but received outputs size is:%d, dy dims is:%d.",
x_grad.size(),
static_cast<size_t>(dy_dim[axis])));
auto vec = vectorize<int64_t>(dy_dim);
vec.erase(vec.begin() + axis);
for (auto& grad : x_grad) {
if (grad) {
grad->set_dims(make_ddim(vec));
grad->set_dtype(out_grad.dtype());
}
}
}
void SwiGLUGradInferMeta(const MetaTensor& x,
const MetaTensor& y,
MetaTensor* x_grad,
MetaTensor* y_grad) {
if (x_grad) {
x_grad->share_meta(x);
}
if (y && y_grad) {
y_grad->share_meta(y);
}
}
void TransposeGradInferMeta(const MetaTensor& x,
const std::vector<int>& axis,
MetaTensor* out) {
size_t x_rank = x.dims().size();
std::vector<int> formatted_axis = axis;
for (size_t i = 0; i < axis.size(); i++) {
if (axis[i] < 0) {
formatted_axis[i] = static_cast<int>(axis[i] + x_rank);
}
}
std::vector<int> reversed_axis(axis);
for (int i = 0; i < static_cast<int>(formatted_axis.size()); i++) {
reversed_axis[formatted_axis[i]] = i;
}
TransposeInferMeta(x, reversed_axis, out);
}
void TransLayoutGradInferMeta(const MetaTensor& x,
const MetaTensor& out_grad,
const std::vector<int>& axis,
MetaTensor* x_grad) {
TransposeGradInferMeta(out_grad, axis, x_grad);
x_grad->set_layout(static_cast<DataLayout>(x.layout()));
}
void UniformRandomInplaceGradInferMeta(const MetaTensor& out_grad,
float min,
float max,
int seed,
int diag_num,
int diag_step,
float diag_val,
MetaTensor* x_grad) {
PADDLE_ENFORCE_NE(
x_grad,
nullptr,
common::errors::InvalidArgument(
"The X@GRAD in UniformRandomInplaceGradInferMeta can't be nullptr."));
auto dims = out_grad.dims();
x_grad->set_dims(dims);
x_grad->set_dtype(out_grad.dtype());
}
void RandomGradInferMeta(const MetaTensor& out_grad, MetaTensor* x_grad) {
PADDLE_ENFORCE_NE(x_grad,
nullptr,
common::errors::InvalidArgument(
"The X@GRAD in RandomGradInferMeta can't be nullptr."));
auto dims = out_grad.dims();
x_grad->set_dims(dims);
x_grad->set_dtype(out_grad.dtype());
}
void UnStackGradInferMeta(const std::vector<const MetaTensor*>& out_grad,
int axis,
MetaTensor* x_grad) {
std::vector<DDim> input_dims(out_grad.size());
for (size_t i = 0; i < out_grad.size(); ++i) {
input_dims[i] = out_grad[i]->dims();
}
for (size_t i = 1; i < input_dims.size(); ++i) {
PADDLE_ENFORCE_EQ(
input_dims[i],
input_dims[0],
common::errors::InvalidArgument(
"The dimensions of all Inputs(Y@GRAD) must be the same,"
"but received Inputs(Y@GRAD)'s %d-th dimension is %d, "
"Inputs(Y@GRAD)'s 0-th to %d-th dimension is %d.",
i,
input_dims[i],
i - 1,
input_dims[0]));
}
int rank = input_dims[0].size();
PADDLE_ENFORCE_GE(axis,
-(rank + 1),
common::errors::InvalidArgument(
"The attribute axis is out of range, it must be "
"inside [-(rank+1), rank+1), where rank = %d",
rank));
PADDLE_ENFORCE_LT(axis,
rank + 1,
common::errors::InvalidArgument(
"The attribute axis is out of range, it must be "
"inside [-(rank+1), rank+1), where rank = %d",
rank));
if (axis < 0) axis += (rank + 1);
auto vec = vectorize<int64_t>(input_dims[0]);
vec.insert(vec.begin() + axis, static_cast<int64_t>(input_dims.size()));
x_grad->set_dims(make_ddim(vec));
x_grad->set_dtype(out_grad[0]->dtype());
}
void WeightOnlyLinearGradInferMeta(const MetaTensor& x,
const MetaTensor& weight,
const MetaTensor& bias,
const MetaTensor& weight_scale,
const MetaTensor& out_grad,
const std::string& weight_dtype,
const int32_t arch,
const int32_t group_size,
MetaTensor* x_grad) {
PADDLE_ENFORCE_EQ(
((arch == 80) || (arch == 86) || (arch == 90) || (arch == 100)),
true,
common::errors::InvalidArgument("Currently weightonly linear grad only "
"support arch = 80, 86, 90 or 100. "));
PADDLE_ENFORCE_EQ(
group_size,
-1,
common::errors::InvalidArgument(
"Currently weightonly linear grad only support per-channel mode. "));
x_grad->set_dims(x.dims());
x_grad->set_dtype(x.dtype());
}
void YoloLossGradInferMeta(const MetaTensor& x,
const MetaTensor& gt_box,
const MetaTensor& gt_label,
const MetaTensor& gt_score,
const MetaTensor& objectness_mask,
const MetaTensor& gt_match_mask,
const MetaTensor& loss_grad,
const std::vector<int>& anchors,
const std::vector<int>& anchor_mask,
int class_num,
float ignore_thresh,
int downsample_ratio,
bool use_label_smooth,
float scale_x_y,
MetaTensor* x_grad,
MetaTensor* gt_box_grad,
MetaTensor* gt_label_grad,
MetaTensor* gt_score_grad) {
if (x_grad) {
x_grad->set_dims(x.dims());
x_grad->set_dtype(x.dtype());
}
}
void IndexAddGradInferMeta(const MetaTensor& index,
const MetaTensor& add_value,
const MetaTensor& out_grad,
int axis,
MetaTensor* x_grad,
MetaTensor* add_value_grad) {
auto do_dims = out_grad.dims();
auto add_value_dims = add_value.dims();
if (x_grad) {
x_grad->set_dims(do_dims);
x_grad->set_dtype(out_grad.dtype());
x_grad->set_layout(out_grad.layout());
x_grad->share_lod(out_grad);
}
if (add_value_grad) {
add_value_grad->set_dims(add_value_dims);
add_value_grad->set_dtype(add_value.dtype());
add_value_grad->set_layout(add_value.layout());
add_value_grad->share_lod(add_value);
}
}
void IndexPutGradInferMeta(const MetaTensor& x,
const std::vector<const MetaTensor*>& indices,
const MetaTensor& value,
const MetaTensor& out_grad,
bool accumulate,
MetaTensor* x_grad,
MetaTensor* value_grad) {
if (x_grad) {
x_grad->share_meta(x);
}
if (value_grad) {
value_grad->share_meta(value);
}
}
void IndexElementwisePutGradInferMeta(
const MetaTensor& x,
const std::vector<const MetaTensor*>& index,
const MetaTensor& out_grad,
const std::vector<int64_t>& input_dims,
const std::vector<int64_t>& input_strides,
const std::vector<int64_t>& index_dims,
const std::vector<int64_t>& index_strides,
const int64_t slice_offset,
MetaTensor* x_grad) {
if (x_grad) {
x_grad->share_meta(x);
}
}
void IndexElementwisePutWithTensorGradInferMeta(
const MetaTensor& x,
const std::vector<const MetaTensor*>& index,
const MetaTensor& value,
const MetaTensor& out_grad,
const std::vector<int64_t>& input_dims,
const std::vector<int64_t>& input_strides,
const std::vector<int64_t>& index_dims,
const std::vector<int64_t>& index_strides,
const int64_t slice_offset,
MetaTensor* x_grad,
MetaTensor* value_grad) {
if (x_grad) {
x_grad->share_meta(x);
}
if (value_grad) {
value_grad->share_meta(value);
}
}
void FusedRopeGradInferMeta(const MetaTensor& sin,
const MetaTensor& cos,
const MetaTensor& position_ids,
const MetaTensor& dout_q,
const MetaTensor& dout_k,
const MetaTensor& dout_v,
bool use_neox_rotary_style,
bool time_major,
float rotary_emb_base,
MetaTensor* dq,
MetaTensor* dk,
MetaTensor* dv) {
auto input_dims = dout_q.dims();
PADDLE_ENFORCE_EQ(input_dims.size(),
4,
common::errors::InvalidArgument(
"Input should be a 4-D tensor of format "
"[batch_size, seq_len, num_heads, head_dim],"
"but got %u.",
input_dims.size()));
if (dout_q && dq) {
dq->set_dims(dout_q.dims());
dq->set_dtype(dout_q.dtype());
}
if (dout_k && dk) {
dk->set_dims(dout_k.dims());
dk->set_dtype(dout_k.dtype());
}
if (dout_v && dv) {
dv->set_dims(dout_v.dims());
dv->set_dtype(dout_v.dtype());
}
}
void SetValueGradInferMeta(const MetaTensor& out_grad,
const MetaTensor& value,
MetaTensor* x_grad,
MetaTensor* value_grad) {
if (x_grad) {
x_grad->set_dims(out_grad.dims());
x_grad->set_dtype(out_grad.dtype());
x_grad->share_lod(out_grad);
}
if (value_grad) {
value_grad->set_dims(value.dims());
value_grad->set_dtype(value.dtype());
value_grad->share_lod(value);
}
}
void CalAuxLossGradInferMeta(const MetaTensor& gate_prob,
const MetaTensor& seqlen_float,
const MetaTensor& ce,
const MetaTensor& l_aux_loss_grad,
const int64_t num_experts,
const bool use_group,
const int64_t moe_k,
MetaTensor* gate_prob_grad) {
auto gate_prob_dims = gate_prob.dims();
PADDLE_ENFORCE_EQ(
gate_prob.dtype(),
l_aux_loss_grad.dtype(),
errors::InvalidArgument(
"The input out_grad type should be equal to gate_prob type"));
gate_prob_grad->set_dims({gate_prob_dims});
gate_prob_grad->set_dtype(gate_prob.dtype());
}
void MoeGateDispatchGradInferMeta(const MetaTensor& combine_weights,
const MetaTensor& scatter_index,
const MetaTensor& expert_id,
const MetaTensor& y_grad,
const MetaTensor& combine_weights_grad,
const int64_t k,
const int64_t capacity,
const bool use_pad,
MetaTensor* x_grad,
MetaTensor* gate_logits_grad) {
auto combine_weights_dims = combine_weights.dims();
auto scatter_index_dims = scatter_index.dims();
auto expert_id_dims = expert_id.dims();
auto y_grad_dims = y_grad.dims();
auto combine_weights_grad_dims = combine_weights_grad.dims();
PADDLE_ENFORCE_EQ(combine_weights_dims.size(),
2,
errors::InvalidArgument(
"Input combine_weights should have 2 dimensions"));
PADDLE_ENFORCE_EQ(
scatter_index_dims.size(),
2,
errors::InvalidArgument("Input scatter_index should have 2 dimensions"));
PADDLE_ENFORCE_EQ(
expert_id_dims.size(),
2,
errors::InvalidArgument("Input expert_id should have 2 dimensions"));
PADDLE_ENFORCE_EQ(
y_grad_dims.size(),
2,
errors::InvalidArgument("Input y_grad should have 2 dimensions"));
PADDLE_ENFORCE_EQ(combine_weights_grad_dims.size(),
2,
errors::InvalidArgument(
"Input combine_weights_grad should have 2 dimensions"));
int64_t num_experts = y_grad_dims[0] / capacity;
int64_t hidden_size = y_grad_dims[1];
int64_t num_rows = scatter_index_dims[1];
x_grad->set_dims(make_ddim({num_rows, hidden_size}));
x_grad->set_dtype(y_grad.dtype());
gate_logits_grad->set_dims(make_ddim({num_rows, num_experts}));
gate_logits_grad->set_dtype(DataType::FLOAT32);
}
void MoeGateDispatchAutoGradInferMeta(const MetaTensor& combine_weights,
const MetaTensor& scatter_index,
const MetaTensor& expert_id,
const MetaTensor& y_grad,
const MetaTensor& combine_weights_grad,
const int64_t k,
const int64_t capacity,
const bool use_pad,
MetaTensor* x_grad,
MetaTensor* gate_logits_grad) {
auto combine_weights_dims = combine_weights.dims();
auto scatter_index_dims = scatter_index.dims();
auto expert_id_dims = expert_id.dims();
auto y_grad_dims = y_grad.dims();
auto combine_weights_grad_dims = combine_weights_grad.dims();
PADDLE_ENFORCE_EQ(combine_weights_dims.size(),
2,
errors::InvalidArgument(
"Input combine_weights should have 2 dimensions"));
PADDLE_ENFORCE_EQ(
scatter_index_dims.size(),
2,
errors::InvalidArgument("Input scatter_index should have 2 dimensions"));
PADDLE_ENFORCE_EQ(
expert_id_dims.size(),
2,
errors::InvalidArgument("Input expert_id should have 2 dimensions"));
PADDLE_ENFORCE_EQ(
y_grad_dims.size(),
3,
errors::InvalidArgument("Input y_grad should have 3 dimensions"));
PADDLE_ENFORCE_EQ(combine_weights_grad_dims.size(),
2,
errors::InvalidArgument(
"Input combine_weights_grad should have 2 dimensions"));
int64_t num_experts = y_grad_dims[0];
int64_t hidden_size = y_grad_dims[2];
int64_t num_rows = scatter_index_dims[1];
x_grad->set_dims(make_ddim({num_rows, hidden_size}));
x_grad->set_dtype(y_grad.dtype());
gate_logits_grad->set_dims(make_ddim({num_rows, num_experts}));
gate_logits_grad->set_dtype(DataType::FLOAT32);
}
void FusedRMSNormGradInferMeta(const MetaTensor& x,
const MetaTensor& scale,
const MetaTensor& invvar,
const MetaTensor& dy,
float epsilon,
MetaTensor* x_grad,
MetaTensor* scale_grad) {
PADDLE_ENFORCE_EQ(
x.dtype() == DataType::FLOAT32 || x.dtype() == DataType::FLOAT16 ||
x.dtype() == DataType::BFLOAT16,
true,
common::errors::InvalidArgument(
"The dtype of x must be FLOAT32, FLOAT16 or BFLOAT16, but got [%s]",
x.dtype()));
PADDLE_ENFORCE_EQ(
scale.dtype() == DataType::FLOAT32 ||
scale.dtype() == DataType::FLOAT16 ||
scale.dtype() == DataType::BFLOAT16,
true,
common::errors::InvalidArgument("The dtype of scale must be FLOAT32, "
"FLOAT16 or BFLOAT16, but got [%s]",
scale.dtype()));
if (x_grad && x) {
x_grad->share_meta(x);
}
if (scale_grad && scale) {
scale_grad->share_meta(scale);
}
}
PADDLE_API void FastLayerNormGradInfermeta(const MetaTensor& x,
const MetaTensor& scale,
const MetaTensor& mean,
const MetaTensor& invvar,
const MetaTensor& y_grad,
float epsilon,
MetaTensor* x_grad,
MetaTensor* scale_grad,
MetaTensor* bias_grad) {
PADDLE_ENFORCE_EQ(
x.dtype() == DataType::FLOAT32 || x.dtype() == DataType::FLOAT16 ||
x.dtype() == DataType::BFLOAT16,
true,
common::errors::InvalidArgument(
"The dtype of x must be FLOAT32, FLOAT16 or BFLOAT16, but got [%s]",
x.dtype()));
PADDLE_ENFORCE_EQ(
scale.dtype() == DataType::FLOAT32 ||
scale.dtype() == DataType::FLOAT16 ||
scale.dtype() == DataType::BFLOAT16,
true,
common::errors::InvalidArgument("The dtype of scale must be FLOAT32, "
"FLOAT16 or BFLOAT16, but got [%s]",
scale.dtype()));
if (x_grad && x) {
x_grad->share_meta(x);
}
if (scale_grad && scale) {
scale_grad->share_meta(scale);
}
if (bias_grad) {
bias_grad->share_meta(scale);
}
}
PADDLE_API void FastRMSNormGradInfermeta(const MetaTensor& x,
const MetaTensor& scale,
const MetaTensor& invvar,
const MetaTensor& y_grad,
float epsilon,
MetaTensor* x_grad,
MetaTensor* scale_grad) {
PADDLE_ENFORCE_EQ(
x.dtype() == DataType::FLOAT32 || x.dtype() == DataType::FLOAT16 ||
x.dtype() == DataType::BFLOAT16,
true,
common::errors::InvalidArgument(
"The dtype of x must be FLOAT32, FLOAT16 or BFLOAT16, but got [%s]",
x.dtype()));
PADDLE_ENFORCE_EQ(
scale.dtype() == DataType::FLOAT32 ||
scale.dtype() == DataType::FLOAT16 ||
scale.dtype() == DataType::BFLOAT16,
true,
common::errors::InvalidArgument("The dtype of scale must be FLOAT32, "
"FLOAT16 or BFLOAT16, but got [%s]",
scale.dtype()));
if (x_grad && x) {
x_grad->share_meta(x);
}
if (scale_grad && scale) {
scale_grad->share_meta(scale);
}
}
void IndexElementwiseGetGradInferMeta(
const MetaTensor& x,
const std::vector<const MetaTensor*>& index,
const MetaTensor& out_grad,
const std::vector<int64_t>& input_dims,
const std::vector<int64_t>& input_strides,
const std::vector<int64_t>& index_dims,
const std::vector<int64_t>& index_strides,
const int64_t slice_offset,
const bool accumulate,
const bool is_combined,
MetaTensor* x_grad) {
if (x_grad) {
x_grad->share_meta(x);
}
}
} // namespace phi