/usr/local/lib64/python3.6/site-packages/torch/include/ATen/native
Edit: /usr/local/lib64/python3.6/site-packages/torch/include/ATen/native/BatchLinearAlgebra.h (8246B)
#pragma once
#include
#include
#include
#include
#include
namespace at { namespace native {
enum class LapackLstsqDriverType : int64_t { Gels, Gelsd, Gelsy, Gelss};
#if AT_BUILD_WITH_LAPACK()
// Define per-batch functions to be used in the implementation of batched
// linear algebra operations
template
void lapackCholesky(char uplo, int n, scalar_t *a, int lda, int *info);
template
void lapackCholeskyInverse(char uplo, int n, scalar_t *a, int lda, int *info);
template
void lapackEig(char jobvl, char jobvr, int n, scalar_t *a, int lda, scalar_t *w, scalar_t* vl, int ldvl, scalar_t *vr, int ldvr, scalar_t *work, int lwork, value_t *rwork, int *info);
template
void lapackGeqrf(int m, int n, scalar_t *a, int lda, scalar_t *tau, scalar_t *work, int lwork, int *info);
template
void lapackOrgqr(int m, int n, int k, scalar_t *a, int lda, scalar_t *tau, scalar_t *work, int lwork, int *info);
template
void lapackOrmqr(char side, char trans, int m, int n, int k, scalar_t *a, int lda, scalar_t *tau, scalar_t *c, int ldc, scalar_t *work, int lwork, int *info);
template
void lapackSyevd(char jobz, char uplo, int n, scalar_t* a, int lda, value_t* w, scalar_t* work, int lwork, value_t* rwork, int lrwork, int* iwork, int liwork, int* info);
template
void lapackGels(char trans, int m, int n, int nrhs,
scalar_t *a, int lda, scalar_t *b, int ldb,
scalar_t *work, int lwork, int *info);
template
void lapackGelsd(int m, int n, int nrhs,
scalar_t *a, int lda, scalar_t *b, int ldb,
value_t *s, value_t rcond, int *rank,
scalar_t* work, int lwork,
value_t *rwork, int* iwork, int *info);
template
void lapackGelsy(int m, int n, int nrhs,
scalar_t *a, int lda, scalar_t *b, int ldb,
int *jpvt, value_t rcond, int *rank,
scalar_t *work, int lwork, value_t* rwork, int *info);
template
void lapackGelss(int m, int n, int nrhs,
scalar_t *a, int lda, scalar_t *b, int ldb,
value_t *s, value_t rcond, int *rank,
scalar_t *work, int lwork,
value_t *rwork, int *info);
template
struct lapackLstsq_impl;
template
struct lapackLstsq_impl {
static void call(
char trans, int m, int n, int nrhs,
scalar_t *a, int lda, scalar_t *b, int ldb,
scalar_t *work, int lwork, int *info, // Gels flavor
int *jpvt, value_t rcond, int *rank, value_t* rwork, // Gelsy flavor
value_t *s, // Gelss flavor
int *iwork // Gelsd flavor
) {
lapackGels(
trans, m, n, nrhs,
a, lda, b, ldb,
work, lwork, info);
}
};
template
struct lapackLstsq_impl {
static void call(
char trans, int m, int n, int nrhs,
scalar_t *a, int lda, scalar_t *b, int ldb,
scalar_t *work, int lwork, int *info, // Gels flavor
int *jpvt, value_t rcond, int *rank, value_t* rwork, // Gelsy flavor
value_t *s, // Gelss flavor
int *iwork // Gelsd flavor
) {
lapackGelsy(
m, n, nrhs,
a, lda, b, ldb,
jpvt, rcond, rank,
work, lwork, rwork, info);
}
};
template
struct lapackLstsq_impl {
static void call(
char trans, int m, int n, int nrhs,
scalar_t *a, int lda, scalar_t *b, int ldb,
scalar_t *work, int lwork, int *info, // Gels flavor
int *jpvt, value_t rcond, int *rank, value_t* rwork, // Gelsy flavor
value_t *s, // Gelss flavor
int *iwork // Gelsd flavor
) {
lapackGelsd(
m, n, nrhs,
a, lda, b, ldb,
s, rcond, rank,
work, lwork,
rwork, iwork, info);
}
};
template
struct lapackLstsq_impl {
static void call(
char trans, int m, int n, int nrhs,
scalar_t *a, int lda, scalar_t *b, int ldb,
scalar_t *work, int lwork, int *info, // Gels flavor
int *jpvt, value_t rcond, int *rank, value_t* rwork, // Gelsy flavor
value_t *s, // Gelss flavor
int *iwork // Gelsd flavor
) {
lapackGelss(
m, n, nrhs,
a, lda, b, ldb,
s, rcond, rank,
work, lwork,
rwork, info);
}
};
template
void lapackLstsq(
char trans, int m, int n, int nrhs,
scalar_t *a, int lda, scalar_t *b, int ldb,
scalar_t *work, int lwork, int *info, // Gels flavor
int *jpvt, value_t rcond, int *rank, value_t* rwork, // Gelsy flavor
value_t *s, // Gelss flavor
int *iwork // Gelsd flavor
) {
lapackLstsq_impl::call(
trans, m, n, nrhs,
a, lda, b, ldb,
work, lwork, info,
jpvt, rcond, rank, rwork,
s,
iwork);
}
template
void lapackLuSolve(char trans, int n, int nrhs, scalar_t *a, int lda, int *ipiv, scalar_t *b, int ldb, int *info);
template
void lapackLu(int m, int n, scalar_t *a, int lda, int *ipiv, int *info);
#endif
#if AT_BUILD_WITH_BLAS()
template
void blasTriangularSolve(char side, char uplo, char trans, char diag, int n, int nrhs, scalar_t* a, int lda, scalar_t* b, int ldb);
#endif
using cholesky_fn = void (*)(const Tensor& /*input*/, const Tensor& /*info*/, bool /*upper*/);
DECLARE_DISPATCH(cholesky_fn, cholesky_stub);
using cholesky_inverse_fn = Tensor& (*)(Tensor& /*result*/, Tensor& /*infos*/, bool /*upper*/);
DECLARE_DISPATCH(cholesky_inverse_fn, cholesky_inverse_stub);
using eig_fn = std::tuple (*)(const Tensor&, bool&);
DECLARE_DISPATCH(eig_fn, eig_stub);
using linalg_eig_fn = void (*)(Tensor& /*eigenvalues*/, Tensor& /*eigenvectors*/, Tensor& /*infos*/, const Tensor& /*input*/, bool /*compute_eigenvectors*/);
DECLARE_DISPATCH(linalg_eig_fn, linalg_eig_stub);
using geqrf_fn = void (*)(const Tensor& /*input*/, const Tensor& /*tau*/);
DECLARE_DISPATCH(geqrf_fn, geqrf_stub);
using orgqr_fn = Tensor& (*)(Tensor& /*result*/, const Tensor& /*tau*/);
DECLARE_DISPATCH(orgqr_fn, orgqr_stub);
using ormqr_fn = void (*)(const Tensor& /*input*/, const Tensor& /*tau*/, const Tensor& /*other*/, bool /*left*/, bool /*transpose*/);
DECLARE_DISPATCH(ormqr_fn, ormqr_stub);
using linalg_eigh_fn = void (*)(
const Tensor& /*eigenvalues*/,
const Tensor& /*eigenvectors*/,
const Tensor& /*infos*/,
bool /*upper*/,
bool /*compute_eigenvectors*/);
DECLARE_DISPATCH(linalg_eigh_fn, linalg_eigh_stub);
using lstsq_fn = void (*)(
const Tensor& /*a*/,
Tensor& /*b*/,
Tensor& /*rank*/,
Tensor& /*singular_values*/,
Tensor& /*infos*/,
double /*rcond*/,
std::string /*driver_name*/);
DECLARE_DISPATCH(lstsq_fn, lstsq_stub);
using triangular_solve_fn = void (*)(
Tensor& /*A*/,
Tensor& /*B*/,
bool /*left*/,
bool /*upper*/,
TransposeType /*transpose*/,
bool /*unitriangular*/);
DECLARE_DISPATCH(triangular_solve_fn, triangular_solve_stub);
using lu_fn = void (*)(
const Tensor& /*input*/,
const Tensor& /*pivots*/,
const Tensor& /*infos*/,
bool /*compute_pivots*/);
DECLARE_DISPATCH(lu_fn, lu_stub);
using lu_solve_fn = void (*)(
const Tensor& /*b*/,
const Tensor& /*lu*/,
const Tensor& /*pivots*/);
DECLARE_DISPATCH(lu_solve_fn, lu_solve_stub);
using lu_solve_trans_fn = void (*)(
const Tensor& /*b*/,
const Tensor& /*lu*/,
const Tensor& /*pivots*/,
TransposeType /*trans*/);
DECLARE_DISPATCH(lu_solve_trans_fn, lu_solve_trans_stub);
}} // namespace at::native