File: transpose_op.h

package info (click to toggle)
pytorch 1.13.1%2Bdfsg-4
  • links: PTS, VCS
  • area: main
  • in suites: bookworm
  • size: 139,252 kB
  • sloc: cpp: 1,100,274; python: 706,454; ansic: 83,052; asm: 7,618; java: 3,273; sh: 2,841; javascript: 612; makefile: 323; xml: 269; ruby: 185; yacc: 144; objc: 68; lex: 44
file content (77 lines) | stat: -rw-r--r-- 2,086 bytes parent folder | download
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
#ifndef CAFFE2_OPERATORS_TRANSPOSE_H_
#define CAFFE2_OPERATORS_TRANSPOSE_H_

#include <algorithm>
#include <vector>

#include "caffe2/core/context.h"
#include "caffe2/core/operator.h"
#include "caffe2/utils/math.h"
#include "c10/util/irange.h"

namespace caffe2 {

template <class Context>
class TransposeOp : public Operator<Context> {
 public:
  USE_OPERATOR_CONTEXT_FUNCTIONS;
  USE_DISPATCH_HELPER;

  template <class... Args>
  explicit TransposeOp(Args&&... args)
      : Operator<Context>(std::forward<Args>(args)...),
        axes_(this->template GetRepeatedArgument<int>("axes")) {
    // We will check the legality of axes_: it should be from 0 to axes_.size().
    std::vector<int> axes_sorted = axes_;
    std::sort(axes_sorted.begin(), axes_sorted.end());
    for (const auto i : c10::irange(axes_sorted.size())) {
      if (axes_sorted[i] != static_cast<int64_t>(i)) {
        CAFFE_THROW("Axes should be a permutation of 0 to ndim.");
      }
    }
  }

  bool RunOnDevice() override {
    // Do the actual transpose, which is implemented in DoRunWithType().
    return DispatchHelper<TensorTypes<float, double, int, int64_t>>::call(
        this, Input(0));
  }

 protected:
  template <typename T>
  void TransposeImpl(const Tensor& X, Tensor* Y) {
    const int ndim = X.dim();
    if (axes_.empty()) {
      axes_.resize(ndim);
      std::iota(axes_.rbegin(), axes_.rend(), 0);
    } else {
      CAFFE_ENFORCE_EQ(ndim, axes_.size());
    }
    const at::IntArrayRef X_dims = X.sizes();
    std::vector<std::int64_t> Y_dims(ndim);
    for (const auto i : c10::irange(ndim)) {
      Y_dims[i] = X_dims[axes_[i]];
    }
    Y->Resize(Y_dims);
    math::Transpose<std::int64_t, T, Context>(
        X_dims.size(),
        X_dims.data(),
        axes_.data(),
        X.template data<T>(),
        Y->template mutable_data<T>(),
        &context_);
  }

 private:
  template <typename T>
  bool DoRunWithType() {
    TransposeImpl<T>(Input(0), Output(0));
    return true;
  }

  std::vector<int> axes_;
};

} // namespace caffe2

#endif // CAFFE2_OPERATORS_TRANSPOSE_H_