chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,343 @@
|
||||
// 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 "paddle/phi/core/dense_tensor.h"
|
||||
#include "paddle/phi/kernels/activation_kernel.h"
|
||||
#include "paddle/phi/kernels/cast_kernel.h"
|
||||
#include "paddle/phi/kernels/complex_kernel.h"
|
||||
#include "paddle/phi/kernels/diag_kernel.h"
|
||||
#include "paddle/phi/kernels/diagonal_kernel.h"
|
||||
#include "paddle/phi/kernels/elementwise_add_kernel.h"
|
||||
#include "paddle/phi/kernels/elementwise_multiply_kernel.h"
|
||||
#include "paddle/phi/kernels/elementwise_subtract_kernel.h"
|
||||
#include "paddle/phi/kernels/funcs/math_function.h"
|
||||
#include "paddle/phi/kernels/matmul_kernel.h"
|
||||
#include "paddle/phi/kernels/slice_kernel.h"
|
||||
#include "paddle/phi/kernels/transpose_kernel.h"
|
||||
|
||||
namespace phi {
|
||||
|
||||
template <class T, class Context>
|
||||
static DenseTensor Fill(const Context& dev_ctx,
|
||||
std::vector<int64_t> shape,
|
||||
T fill_value) {
|
||||
DenseTensor ret;
|
||||
ret.Resize(shape);
|
||||
dev_ctx.template Alloc<T>(&ret);
|
||||
funcs::SetConstant<Context, T>()(dev_ctx, &ret, fill_value);
|
||||
return ret;
|
||||
}
|
||||
|
||||
template <class T, class Context>
|
||||
static DenseTensor Eye(const Context& dev_ctx, int64_t n) {
|
||||
auto output = Fill<T, Context>(dev_ctx, {n}, T(1));
|
||||
auto ret = Diag<T, Context>(dev_ctx, output, 0, 0);
|
||||
return ret;
|
||||
}
|
||||
|
||||
template <class T, class Context>
|
||||
static DenseTensor Infinits(const Context& dev_ctx,
|
||||
std::vector<int64_t> shape) {
|
||||
auto value = static_cast<T>(std::numeric_limits<double>::infinity());
|
||||
return Fill<T, Context>(dev_ctx, shape, value);
|
||||
}
|
||||
|
||||
static DenseTensor Unsqueeze(const DenseTensor& x, int axis = 0) {
|
||||
// don't copy data, only change the dims
|
||||
DenseTensor out;
|
||||
out.ShareDataWith(x);
|
||||
std::vector<int64_t> out_shape = vectorize<int64_t>(x.dims());
|
||||
if (axis >= 0) {
|
||||
auto index = (out_shape.begin() + axis);
|
||||
out_shape.insert(index, 1);
|
||||
} else if (axis < 0) {
|
||||
auto index = (out_shape.end() + axis + 1);
|
||||
out_shape.insert(index, 1);
|
||||
}
|
||||
out.Resize(out_shape);
|
||||
return out;
|
||||
}
|
||||
|
||||
template <typename T, typename Context>
|
||||
DenseTensor Hermitian(const Context& dev_ctx, const DenseTensor& x) {
|
||||
return TransposeLast2Dim<T>(dev_ctx, Conj<T, Context>(dev_ctx, x));
|
||||
}
|
||||
|
||||
template <typename T, typename Context>
|
||||
struct SvdGradFunctor {
|
||||
void operator()(const Context& dev_ctx,
|
||||
const DenseTensor& u,
|
||||
const DenseTensor& vh,
|
||||
const DenseTensor& s,
|
||||
const optional<DenseTensor>& u_grad,
|
||||
const optional<DenseTensor>& vh_grad,
|
||||
const optional<DenseTensor>& s_grad,
|
||||
bool full_matrices,
|
||||
DenseTensor* x_grad) {
|
||||
const auto& dX = *x_grad;
|
||||
int64_t m = dX.dims()[dX.dims().size() - 2];
|
||||
int64_t n = dX.dims()[dX.dims().size() - 1];
|
||||
int64_t k = s.dims()[s.dims().size() - 1];
|
||||
DenseTensor U, VH, dU, dV, dVH;
|
||||
if (full_matrices) {
|
||||
// if full_matrices is set, slice the U and VT to k columns
|
||||
U = Slice<T, Context>(dev_ctx, u, {u.dims().size() - 1}, {0}, {k});
|
||||
// If m < n for input matrices A, we partition A = [X|Y] and R = [U|V]
|
||||
|
||||
VH = Slice<T, Context>(dev_ctx, vh, {vh.dims().size() - 2}, {0}, {k});
|
||||
if (u_grad.get_ptr() != nullptr) {
|
||||
dU = Slice<T, Context>(
|
||||
dev_ctx, *(u_grad.get_ptr()), {u.dims().size() - 1}, {0}, {k});
|
||||
}
|
||||
if (vh_grad.get_ptr() != nullptr) {
|
||||
dVH = Slice<T, Context>(
|
||||
dev_ctx, *(vh_grad.get_ptr()), {vh.dims().size() - 2}, {0}, {k});
|
||||
}
|
||||
} else {
|
||||
U = u;
|
||||
VH = vh;
|
||||
if (u_grad.get_ptr() != nullptr) {
|
||||
dU = *(u_grad.get_ptr());
|
||||
}
|
||||
if (vh_grad.get_ptr() != nullptr) {
|
||||
dVH = *(vh_grad.get_ptr());
|
||||
}
|
||||
}
|
||||
auto s_inverse = Pow<T, Context>(dev_ctx, s, -1);
|
||||
auto s_square = Pow<T, Context>(dev_ctx, s, 2);
|
||||
auto F = Subtract<T, Context>(
|
||||
dev_ctx, Unsqueeze(s_square, -2), Unsqueeze(s_square, -1));
|
||||
F = Add<T, Context>(
|
||||
dev_ctx,
|
||||
F,
|
||||
Diag<T, Context>(dev_ctx, Infinits<T, Context>(dev_ctx, {k}), 0, 0));
|
||||
F = Pow<T, Context>(dev_ctx, F, -1);
|
||||
DenseTensor sigma_term = Fill<T, Context>(dev_ctx, {1}, T(0.0));
|
||||
DenseTensor u_term = Fill<T, Context>(dev_ctx, {1}, T(0.0));
|
||||
DenseTensor v_term = Fill<T, Context>(dev_ctx, {1}, T(0.0));
|
||||
|
||||
if (s_grad.get_ptr() != nullptr) {
|
||||
const DenseTensor& gS = *(s_grad.get_ptr());
|
||||
sigma_term = Multiply<T, Context>(dev_ctx, Unsqueeze(gS, -2), U);
|
||||
sigma_term = Matmul<T, Context>(dev_ctx, sigma_term, VH);
|
||||
}
|
||||
|
||||
if (u_grad.get_ptr() != nullptr) {
|
||||
auto UTG = Matmul<T, Context>(dev_ctx, U, dU, true, false);
|
||||
auto GTU = Matmul<T, Context>(dev_ctx, dU, U, true, false);
|
||||
u_term = Multiply<T, Context>(
|
||||
dev_ctx,
|
||||
Multiply<T, Context>(
|
||||
dev_ctx, Subtract<T, Context>(dev_ctx, UTG, GTU), F),
|
||||
Unsqueeze(s, -2));
|
||||
u_term = Matmul<T, Context>(dev_ctx, U, u_term);
|
||||
if (m > k) {
|
||||
auto project = Subtract<T, Context>(
|
||||
dev_ctx,
|
||||
Eye<T, Context>(dev_ctx, m),
|
||||
Matmul<T, Context>(dev_ctx, U, U, false, true));
|
||||
u_term = Add<T, Context>(
|
||||
dev_ctx,
|
||||
u_term,
|
||||
Multiply<T, Context>(dev_ctx,
|
||||
Matmul<T, Context>(dev_ctx, project, dU),
|
||||
Unsqueeze(s_inverse, -2)));
|
||||
}
|
||||
u_term = Matmul<T, Context>(dev_ctx, u_term, VH);
|
||||
}
|
||||
if (vh_grad.get_ptr() != nullptr) {
|
||||
auto UTG = Matmul<T, Context>(dev_ctx, VH, dVH, false, true);
|
||||
auto GTU = Matmul<T, Context>(dev_ctx, dVH, VH, false, true);
|
||||
v_term = Multiply<T, Context>(
|
||||
dev_ctx,
|
||||
Matmul<T, Context>(
|
||||
dev_ctx,
|
||||
Multiply<T, Context>(
|
||||
dev_ctx, Subtract<T, Context>(dev_ctx, UTG, GTU), F),
|
||||
VH),
|
||||
Unsqueeze(s, -1));
|
||||
if (n > k) {
|
||||
auto project = Subtract<T, Context>(
|
||||
dev_ctx,
|
||||
Eye<T, Context>(dev_ctx, n),
|
||||
Matmul<T, Context>(dev_ctx, VH, VH, true, false));
|
||||
v_term = Add<T, Context>(
|
||||
dev_ctx,
|
||||
v_term,
|
||||
Multiply<T, Context>(dev_ctx,
|
||||
Matmul<T, Context>(dev_ctx, dVH, project),
|
||||
Unsqueeze(s_inverse, -1)));
|
||||
}
|
||||
v_term = Matmul<T, Context>(dev_ctx, U, v_term);
|
||||
}
|
||||
|
||||
*x_grad = Add<T, Context>(
|
||||
dev_ctx, Add<T, Context>(dev_ctx, u_term, sigma_term), v_term);
|
||||
}
|
||||
};
|
||||
|
||||
template <typename T, typename Context>
|
||||
struct SvdGradFunctor<dtype::complex<T>, Context> {
|
||||
void operator()(const Context& dev_ctx,
|
||||
const DenseTensor& u,
|
||||
const DenseTensor& vh,
|
||||
const DenseTensor& s,
|
||||
const optional<DenseTensor>& u_grad,
|
||||
const optional<DenseTensor>& vh_grad,
|
||||
const optional<DenseTensor>& s_grad,
|
||||
bool full_matrices,
|
||||
DenseTensor* x_grad) {
|
||||
using C = dtype::complex<T>;
|
||||
const auto& dX = *x_grad;
|
||||
int64_t m = dX.dims()[dX.dims().size() - 2];
|
||||
int64_t n = dX.dims()[dX.dims().size() - 1];
|
||||
int64_t k = s.dims()[s.dims().size() - 1];
|
||||
DenseTensor S = Cast<T, Context>(dev_ctx, s, u.dtype());
|
||||
DenseTensor U, VH, dU, dV, dVH;
|
||||
|
||||
if (full_matrices) {
|
||||
// if full_matrices is set, slice the U and VT to k columns
|
||||
U = Slice<C, Context>(dev_ctx, u, {u.dims().size() - 1}, {0}, {k});
|
||||
// If m < n for input matrices A, we partition A = [X|Y] and R = [U|V]
|
||||
VH = Slice<C, Context>(dev_ctx, vh, {vh.dims().size() - 2}, {0}, {k});
|
||||
if (u_grad.get_ptr() != nullptr) {
|
||||
dU = Slice<C, Context>(
|
||||
dev_ctx, *(u_grad.get_ptr()), {u.dims().size() - 1}, {0}, {k});
|
||||
}
|
||||
if (vh_grad.get_ptr() != nullptr) {
|
||||
dVH = Slice<C, Context>(
|
||||
dev_ctx, *(vh_grad.get_ptr()), {vh.dims().size() - 2}, {0}, {k});
|
||||
}
|
||||
} else {
|
||||
U = u;
|
||||
VH = vh;
|
||||
if (u_grad.get_ptr() != nullptr) {
|
||||
dU = *(u_grad.get_ptr());
|
||||
}
|
||||
if (vh_grad.get_ptr() != nullptr) {
|
||||
dVH = *(vh_grad.get_ptr());
|
||||
}
|
||||
}
|
||||
auto s_inverse = Pow<C, Context>(dev_ctx, S, -1);
|
||||
auto s_square = Pow<C, Context>(dev_ctx, S, 2);
|
||||
auto F = Subtract<C, Context>(
|
||||
dev_ctx, Unsqueeze(s_square, -2), Unsqueeze(s_square, -1));
|
||||
F = Add<C, Context>(
|
||||
dev_ctx,
|
||||
F,
|
||||
Diag<C, Context>(dev_ctx, Infinits<C, Context>(dev_ctx, {k}), 0, 0));
|
||||
F = Pow<C, Context>(dev_ctx, F, -1);
|
||||
DenseTensor sigma_term = Fill<C, Context>(dev_ctx, {1}, C(0.0));
|
||||
DenseTensor u_term = Fill<C, Context>(dev_ctx, {1}, C(0.0));
|
||||
DenseTensor v_term = Fill<C, Context>(dev_ctx, {1}, C(0.0));
|
||||
DenseTensor extra = Fill<C, Context>(dev_ctx, {1}, C(0.0));
|
||||
|
||||
if (s_grad.get_ptr() != nullptr) {
|
||||
const DenseTensor& gS = *(s_grad.get_ptr());
|
||||
DenseTensor dS = Cast<T, Context>(dev_ctx, gS, u.dtype());
|
||||
sigma_term =
|
||||
Multiply<C, Context>(dev_ctx, Eye<C, Context>(dev_ctx, k), dS);
|
||||
sigma_term = Matmul<C, Context>(dev_ctx, U, sigma_term);
|
||||
sigma_term = Matmul<C, Context>(dev_ctx, sigma_term, VH);
|
||||
}
|
||||
|
||||
const auto skew = [](const Context& dev_ctx, const DenseTensor& A) {
|
||||
return Subtract<C, Context>(
|
||||
dev_ctx, A, Hermitian<C, Context>(dev_ctx, A));
|
||||
};
|
||||
|
||||
if (u_grad.get_ptr() != nullptr) {
|
||||
auto UhgU = skew(
|
||||
dev_ctx,
|
||||
Matmul<C, Context>(dev_ctx, Hermitian<C, Context>(dev_ctx, U), dU));
|
||||
u_term = Multiply<C, Context>(
|
||||
dev_ctx, Multiply<C, Context>(dev_ctx, UhgU, F), Unsqueeze(S, -2));
|
||||
u_term = Matmul<C, Context>(dev_ctx, U, u_term);
|
||||
if (m > k) {
|
||||
auto project = Subtract<C, Context>(
|
||||
dev_ctx,
|
||||
Eye<C, Context>(dev_ctx, m),
|
||||
Matmul<C, Context>(dev_ctx, U, Hermitian<C, Context>(dev_ctx, U)));
|
||||
u_term = Add<C, Context>(
|
||||
dev_ctx,
|
||||
u_term,
|
||||
Multiply<C, Context>(dev_ctx,
|
||||
Matmul<C, Context>(dev_ctx, project, dU),
|
||||
Unsqueeze(s_inverse, -2)));
|
||||
}
|
||||
u_term = Matmul<C, Context>(dev_ctx, u_term, VH);
|
||||
|
||||
// complex extra
|
||||
size_t rank = UhgU.dims().size();
|
||||
extra = Matmul<C, Context>(
|
||||
dev_ctx,
|
||||
Diagonal<C, Context>(dev_ctx, UhgU, 0, rank - 2, rank - 1),
|
||||
Pow<C, Context>(dev_ctx,
|
||||
Multiply<C, Context>(
|
||||
dev_ctx, Fill<C, Context>(dev_ctx, {1}, C(2)), S),
|
||||
-1));
|
||||
extra = Multiply<C, Context>(dev_ctx, Eye<C, Context>(dev_ctx, k), extra);
|
||||
extra = Matmul<C, Context>(dev_ctx, U, extra);
|
||||
extra = Matmul<C, Context>(dev_ctx, extra, VH);
|
||||
}
|
||||
|
||||
if (vh_grad.get_ptr() != nullptr) {
|
||||
auto VhgV = skew(
|
||||
dev_ctx,
|
||||
Matmul<C, Context>(dev_ctx, VH, Hermitian<C, Context>(dev_ctx, dVH)));
|
||||
|
||||
v_term = Multiply<C, Context>(
|
||||
dev_ctx, Unsqueeze(S, -1), Multiply<C, Context>(dev_ctx, VhgV, F));
|
||||
v_term = Matmul<C, Context>(dev_ctx, v_term, VH);
|
||||
if (n > k) {
|
||||
auto project = Subtract<C, Context>(
|
||||
dev_ctx,
|
||||
Eye<C, Context>(dev_ctx, n),
|
||||
Matmul<C, Context>(
|
||||
dev_ctx, Hermitian<C, Context>(dev_ctx, VH), VH));
|
||||
v_term = Add<C, Context>(
|
||||
dev_ctx,
|
||||
v_term,
|
||||
Multiply<C, Context>(dev_ctx,
|
||||
Matmul<C, Context>(dev_ctx, dVH, project),
|
||||
Unsqueeze(s_inverse, -1)));
|
||||
}
|
||||
v_term = Matmul<C, Context>(dev_ctx, U, v_term);
|
||||
}
|
||||
|
||||
*x_grad = Add<C, Context>(
|
||||
dev_ctx,
|
||||
Add<C, Context>(
|
||||
dev_ctx, Add<C, Context>(dev_ctx, u_term, sigma_term), v_term),
|
||||
extra);
|
||||
}
|
||||
};
|
||||
|
||||
template <typename T, typename Context>
|
||||
void SvdGradKernel(const Context& dev_ctx,
|
||||
const DenseTensor& x UNUSED,
|
||||
const DenseTensor& u,
|
||||
const DenseTensor& vh,
|
||||
const DenseTensor& s,
|
||||
const optional<DenseTensor>& u_grad,
|
||||
const optional<DenseTensor>& vh_grad,
|
||||
const optional<DenseTensor>& s_grad,
|
||||
bool full_matrices,
|
||||
DenseTensor* x_grad) {
|
||||
SvdGradFunctor<T, Context>()(
|
||||
dev_ctx, u, vh, s, u_grad, vh_grad, s_grad, full_matrices, x_grad);
|
||||
}
|
||||
|
||||
} // namespace phi
|
||||
Reference in New Issue
Block a user