chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,291 @@
|
||||
// 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.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "glog/logging.h"
|
||||
|
||||
#include "paddle/phi/core/tensor_utils.h"
|
||||
#include "paddle/phi/core/utils/data_type.h"
|
||||
#include "paddle/phi/kernels/complex_kernel.h"
|
||||
#include "paddle/phi/kernels/elementwise_add_kernel.h"
|
||||
#include "paddle/phi/kernels/elementwise_multiply_kernel.h"
|
||||
#include "paddle/phi/kernels/full_kernel.h"
|
||||
#include "paddle/phi/kernels/funcs/complex_functors.h"
|
||||
#include "paddle/phi/kernels/funcs/math_function.h"
|
||||
#include "paddle/phi/kernels/funcs/matrix_inverse.h"
|
||||
#include "paddle/phi/kernels/funcs/unsqueeze.h"
|
||||
#include "paddle/phi/kernels/impl/determinant_grad_kernel_impl.h"
|
||||
#include "paddle/phi/kernels/impl/isfinite_kernel_impl.h"
|
||||
#include "paddle/phi/kernels/slogdeterminant_grad_kernel.h"
|
||||
#include "paddle/phi/kernels/transpose_kernel.h"
|
||||
|
||||
namespace phi {
|
||||
|
||||
template <typename T, typename Context>
|
||||
void SlogDeterminantGradKernel(const Context& dev_ctx,
|
||||
const DenseTensor& x,
|
||||
const DenseTensor& out,
|
||||
const DenseTensor& out_grad,
|
||||
DenseTensor* x_grad) {
|
||||
if (x_grad && x_grad->numel() == 0) {
|
||||
dev_ctx.template Alloc<T>(x_grad);
|
||||
return;
|
||||
}
|
||||
PADDLE_ENFORCE_EQ(
|
||||
out_grad.dims()[0],
|
||||
2,
|
||||
errors::InvalidArgument("The grad tensor of SlogDet should contain two"
|
||||
" grad: sign and absslogdet, but here %ld.",
|
||||
out_grad.dims()[0]));
|
||||
if (x.dims().size() > 2) {
|
||||
PADDLE_ENFORCE_EQ(
|
||||
out_grad.dims().size() + 1,
|
||||
x.dims().size(),
|
||||
errors::InvalidArgument(
|
||||
"The grad tensor of slogdet dims size should 1 less than"
|
||||
" input tensor's, but here differ %d",
|
||||
x.dims().size() - out_grad.dims().size()));
|
||||
}
|
||||
|
||||
// Check Whether the matrix is invertible
|
||||
// (matrix A not invertible) == (absslogdet(A)=0)
|
||||
auto slogdet_vec = out.Split(1, 0);
|
||||
auto absslogdet_val = slogdet_vec[0];
|
||||
if (!detail::CheckMatrixInvertible<T, Context>(dev_ctx, &absslogdet_val)) {
|
||||
// The matrix is not invertible
|
||||
VLOG(3) << "The input matrix not invertible!";
|
||||
x_grad->Resize(x.dims());
|
||||
Full<T>(dev_ctx,
|
||||
vectorize(x.dims()),
|
||||
std::numeric_limits<T>::quiet_NaN(),
|
||||
x_grad);
|
||||
return;
|
||||
}
|
||||
|
||||
// The matrix is invertible
|
||||
// let sl|A| = SlogDeterminant(A)
|
||||
// Ref to https://people.maths.ox.ac.uk/gilesm/files/NA-08-01.pdf
|
||||
// we set dsl|A| = unsqueeze(dslA, [-1, -2]) *
|
||||
// inverse(A).conj().transpose(-2, -1)
|
||||
|
||||
// First: inverse(A)
|
||||
DenseTensor inverse_A;
|
||||
// A must be square matrices!
|
||||
inverse_A.Resize(x.dims());
|
||||
dev_ctx.template Alloc<T>(&inverse_A);
|
||||
|
||||
const auto& mat_dims = x.dims();
|
||||
const int rank = mat_dims.size();
|
||||
int n = mat_dims[rank - 1];
|
||||
int64_t total_batch_size = rank > 2 ? x.numel() / (n * n) : 1;
|
||||
|
||||
// Divide the batch into chunks because of cublasMatInv limitation
|
||||
if (total_batch_size <= 65536) {
|
||||
funcs::MatrixInverseFunctor<Context, T> mat_inv;
|
||||
mat_inv(dev_ctx, x, &inverse_A);
|
||||
} else {
|
||||
constexpr int64_t max_batch_size = 65536;
|
||||
int64_t processed = 0;
|
||||
|
||||
VLOG(3) << "Large batch size detected (" << total_batch_size
|
||||
<< "), processing in chunks of " << max_batch_size;
|
||||
|
||||
while (processed < total_batch_size) {
|
||||
int64_t current_batch =
|
||||
std::min(max_batch_size, total_batch_size - processed);
|
||||
|
||||
// Extract current batch data
|
||||
DenseTensor x_batch;
|
||||
x_batch.ShareDataWith(x);
|
||||
x_batch.Resize({total_batch_size, n, n});
|
||||
x_batch = x_batch.Slice(processed, processed + current_batch);
|
||||
x_batch.Resize({current_batch, n, n});
|
||||
|
||||
DenseTensor inverse_batch;
|
||||
inverse_batch.Resize({current_batch, n, n});
|
||||
dev_ctx.template Alloc<T>(&inverse_batch);
|
||||
|
||||
// Compute the inverse matrix for the current batch
|
||||
funcs::MatrixInverseFunctor<Context, T> mat_inv;
|
||||
mat_inv(dev_ctx, x_batch, &inverse_batch);
|
||||
|
||||
// Copy the result to the output tensor
|
||||
DenseTensor output_slice;
|
||||
output_slice.ShareDataWith(inverse_A);
|
||||
output_slice.Resize({total_batch_size, n, n});
|
||||
output_slice = output_slice.Slice(processed, processed + current_batch);
|
||||
output_slice.Resize({current_batch, n, n});
|
||||
|
||||
Copy(dev_ctx, inverse_batch, dev_ctx.GetPlace(), false, &output_slice);
|
||||
|
||||
processed += current_batch;
|
||||
}
|
||||
}
|
||||
|
||||
VLOG(3) << "inverse(A) dims: " << inverse_A.dims();
|
||||
|
||||
// Second: inverse(A).conj()
|
||||
auto conj_inverse_A = Conj<T>(dev_ctx, inverse_A);
|
||||
|
||||
VLOG(3) << "inverse(A).conj() dims: " << conj_inverse_A.dims();
|
||||
|
||||
// Third: inverse(A).conj().transpose(-2, -1)
|
||||
DenseTensor transpose_inverse_A =
|
||||
TransposeLast2Dim<T>(dev_ctx, conj_inverse_A);
|
||||
VLOG(3) << "inverse(A).conj().transpose(-2, -1) dims: "
|
||||
<< transpose_inverse_A.dims();
|
||||
|
||||
// Fourth: split grad value to [sign_grad, absslogdet_grad]
|
||||
auto grad_vec = out_grad.Split(1, 0);
|
||||
auto det_grad = grad_vec[1];
|
||||
|
||||
// remove useless first dimension
|
||||
int det_grad_size = det_grad.dims().size();
|
||||
std::vector<int64_t> det_grad_vec;
|
||||
for (int64_t i = 1; i < det_grad_size; ++i) {
|
||||
det_grad_vec.emplace_back(det_grad.dims()[i]);
|
||||
}
|
||||
det_grad.Resize(det_grad.dims().reshape(det_grad_vec));
|
||||
|
||||
// Fifth: unsqueeze(dslA, [-1, -2])
|
||||
auto unsqueeze1 = funcs::Unsqueeze(det_grad, -1);
|
||||
auto unsqueeze2 = funcs::Unsqueeze(unsqueeze1, -2);
|
||||
VLOG(3) << "unsqueezed(dslA, [-1, -2]) dims: " << unsqueeze2.dims();
|
||||
|
||||
// Finally: unsqueeze(dslA) * inverse(A)
|
||||
auto res = Multiply<T>(dev_ctx, unsqueeze2, transpose_inverse_A);
|
||||
VLOG(3) << "unsqueeze(dslA) * inverse(A) dims: " << res.dims();
|
||||
|
||||
Copy(dev_ctx, res, dev_ctx.GetPlace(), false, x_grad);
|
||||
x_grad->Resize(x.dims());
|
||||
VLOG(3) << "dsl|A| dims: " << x_grad->dims();
|
||||
}
|
||||
|
||||
template <typename T, typename Context>
|
||||
void SlogDeterminantV2GradKernel(const Context& dev_ctx,
|
||||
const DenseTensor& x,
|
||||
const DenseTensor& sign,
|
||||
const DenseTensor& logdet,
|
||||
const DenseTensor& sign_grad UNUSED,
|
||||
const DenseTensor& logdet_grad,
|
||||
DenseTensor* x_grad) {
|
||||
using RealT = typename dtype::Real<T>;
|
||||
const auto& x_dims = x.dims();
|
||||
const auto& grad_dims = logdet_grad.dims();
|
||||
int x_rank = x_dims.size();
|
||||
int grad_rank = grad_dims.size();
|
||||
|
||||
PADDLE_ENFORCE_GE(
|
||||
x_rank,
|
||||
2,
|
||||
common::errors::InvalidArgument(
|
||||
"Input tensor X's rank must be at least 2, but received %d.",
|
||||
x_rank));
|
||||
|
||||
if (x_rank == 2)
|
||||
PADDLE_ENFORCE_EQ(
|
||||
grad_rank,
|
||||
0,
|
||||
common::errors::InvalidArgument(
|
||||
"For a 2D input tensor X, the gradient tensor (logdet_grad) "
|
||||
"should be a 0D tensor (scalar), but received rank %d.",
|
||||
grad_rank));
|
||||
else if (x_rank > 2)
|
||||
PADDLE_ENFORCE_EQ(
|
||||
grad_rank + 2,
|
||||
x_rank,
|
||||
common::errors::InvalidArgument(
|
||||
"The rank of gradient tensor (logdet_grad) should be 2 less than "
|
||||
"the input tensor X's rank, but received grad rank %d and X rank "
|
||||
"%d.",
|
||||
grad_rank,
|
||||
x_rank));
|
||||
|
||||
dev_ctx.template Alloc<T>(x_grad);
|
||||
if (x_grad->numel() == 0) {
|
||||
return;
|
||||
}
|
||||
|
||||
// Check Whether the matrix is invertible
|
||||
// (matrix A not invertible) == (absslogdet(A)=0)
|
||||
if (!detail::CheckMatrixInvertible<RealT, Context>(dev_ctx, &logdet)) {
|
||||
// The matrix is not invertible
|
||||
VLOG(3) << "The input matrix not invertible!";
|
||||
Full<T>(dev_ctx,
|
||||
vectorize(x.dims()),
|
||||
std::numeric_limits<T>::quiet_NaN(),
|
||||
x_grad);
|
||||
return;
|
||||
}
|
||||
|
||||
// The matrix is invertible
|
||||
// let sl|A| = SlogDeterminant(A)
|
||||
// Ref to https://people.maths.ox.ac.uk/gilesm/files/NA-08-01.pdf
|
||||
// we set dsl|A| = unsqueeze(dslA, [-1, -2]) *
|
||||
// inverse(A).conj().transpose(-2, -1)
|
||||
|
||||
// First: inverse(A)
|
||||
DenseTensor inverse_A;
|
||||
// A must be square matrices!
|
||||
inverse_A.Resize(x_dims);
|
||||
dev_ctx.template Alloc<T>(&inverse_A);
|
||||
|
||||
funcs::MatrixInverseFunctor<Context, T> mat_inv;
|
||||
mat_inv(dev_ctx, x, &inverse_A);
|
||||
|
||||
VLOG(3) << "inverse(A) dims: " << inverse_A.dims();
|
||||
|
||||
// Second: inverse(A).conj() for complex
|
||||
DenseTensor conj_inverse_A;
|
||||
if constexpr (is_complex64_or_complex128<T>::value) {
|
||||
conj_inverse_A = Conj<T>(dev_ctx, inverse_A);
|
||||
VLOG(3) << "Performed complex conjugate.";
|
||||
} else {
|
||||
conj_inverse_A.ShareDataWith(inverse_A);
|
||||
VLOG(3) << "Skipped complex conjugate for real type.";
|
||||
}
|
||||
|
||||
VLOG(3) << "inverse(A).conj() dims: " << conj_inverse_A.dims();
|
||||
|
||||
// Third: inverse(A).conj().transpose(-2, -1)
|
||||
DenseTensor transpose_inverse_A =
|
||||
TransposeLast2Dim<T>(dev_ctx, conj_inverse_A);
|
||||
VLOG(3) << "inverse(A).conj().transpose(-2, -1) dims: "
|
||||
<< transpose_inverse_A.dims();
|
||||
|
||||
DenseTensor logdet_grad_term = logdet_grad;
|
||||
if constexpr (is_complex64_or_complex128<T>::value) {
|
||||
// change logdet_grad datatype from <RealT> to <ComplexT>
|
||||
DenseTensor logdet_grad_complex = Empty<T>(dev_ctx, vectorize(grad_dims));
|
||||
|
||||
int64_t logdet_numel = logdet_grad.numel();
|
||||
funcs::ForRange<Context> for_range(dev_ctx, logdet_numel);
|
||||
funcs::RealToComplexFunctor<T> functor(
|
||||
logdet_grad.data<RealT>(), logdet_grad_complex.data<T>(), logdet_numel);
|
||||
|
||||
for_range(functor);
|
||||
logdet_grad_term = logdet_grad_complex;
|
||||
}
|
||||
DenseTensor unsqueezed_combined_grad = funcs::Unsqueeze(logdet_grad_term, -1);
|
||||
unsqueezed_combined_grad = funcs::Unsqueeze(unsqueezed_combined_grad, -2);
|
||||
VLOG(3) << "unsqueezed_combined_grad dims: "
|
||||
<< unsqueezed_combined_grad.dims();
|
||||
|
||||
Multiply<T, Context>(
|
||||
dev_ctx, unsqueezed_combined_grad, transpose_inverse_A, x_grad);
|
||||
VLOG(3) << x_grad->dims();
|
||||
}
|
||||
|
||||
} // namespace phi
|
||||
Reference in New Issue
Block a user