mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
e04c9195b7
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/17670 Update math::Transpose to support tensor with size > 2G i-am-not-moving-c2-to-c10 Differential Revision: D14313624 fbshipit-source-id: 0b4a85b913972e5a8981f0d40d0c539407b98f30
32 lines
788 B
C++
32 lines
788 B
C++
#ifndef CAFFE2_UTILS_MATH_TRANSPOSE_H_
|
|
#define CAFFE2_UTILS_MATH_TRANSPOSE_H_
|
|
|
|
#include "caffe2/core/common.h"
|
|
#include "caffe2/core/types.h"
|
|
|
|
namespace caffe2 {
|
|
namespace math {
|
|
|
|
// Transpose tensor X with dims by axes and write the result to tensor Y.
|
|
template <typename TIndex, typename TData, class Context>
|
|
CAFFE2_API void Transpose(
|
|
int ndim,
|
|
const TIndex* dims,
|
|
const int* axes,
|
|
const TData* X,
|
|
TData* Y,
|
|
Context* context);
|
|
|
|
template <typename T, class Context>
|
|
CAFFE2_API void
|
|
NCHW2NHWC(int N, int C, int HxW, const T* X, T* Y, Context* context);
|
|
|
|
template <typename T, class Context>
|
|
CAFFE2_API void
|
|
NHWC2NCHW(int N, int C, int HxW, const T* X, T* Y, Context* context);
|
|
|
|
} // namespace math
|
|
} // namespace caffe2
|
|
|
|
#endif // CAFFE2_UTILS_MATH_TRANSPOSE_H_
|