int_array.h 3.3 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16
/* Copyright (c) 2021 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

17 18
#include "paddle/phi/api/ext/exception.h"
#include "paddle/phi/api/include/tensor.h"
19 20 21 22 23

namespace paddle {
namespace experimental {

template <typename T>
24
class IntArrayBase {
25 26
 public:
  // Constructor support implicit
27
  IntArrayBase() = default;
28

29
  IntArrayBase(const std::vector<int64_t>& vec) : array_(vec) {}  // NOLINT
30

31
  IntArrayBase(const std::vector<int32_t>& vec) {  // NOLINT
32 33 34
    array_.insert(array_.begin(), vec.begin(), vec.end());
  }

35
  IntArrayBase(std::initializer_list<int64_t> array_list)
36 37
      : array_(array_list) {}

38 39
  IntArrayBase(const int64_t* data_value, int64_t n) {
    AssignData(data_value, n);
40 41
  }

42 43
  IntArrayBase(const int32_t* data_value, int64_t n) {
    AssignData(data_value, n);
44 45
  }

46
  bool FromTensor() const { return is_from_tensor_; }
C
chentianyu03 已提交
47

48
  void SetFromTensor(bool val) { is_from_tensor_ = val; }
C
chentianyu03 已提交
49

50
  // The Tensor must have one dim
51
  IntArrayBase(const T& tensor);  // NOLINT
52 53

  // The Tensor in vec must have only one element
54
  IntArrayBase(const std::vector<T>& tensor_list);  // NOLINT
55 56

  template <typename OtherT>
57
  IntArrayBase(const IntArrayBase<OtherT>& other) : array_(other.GetData()) {}
58

59 60
  size_t size() const { return array_.size(); }

Z
zyfncg 已提交
61 62
  int64_t operator[](int64_t i) const { return array_[i]; }

63 64 65 66 67 68
  const std::vector<int64_t>& GetData() const { return array_; }

 private:
  /// \brief Assign the data_ from const data pointer value of type T.
  template <typename TYPE>
  void AssignData(const TYPE* value_data, int64_t n) {
69
    if (value_data || n == 0) {
70 71 72 73 74 75 76 77 78
      array_.reserve(n);
      for (auto i = 0; i < n; ++i) {
        array_.push_back(static_cast<int64_t>(value_data[i]));
      }
    } else {
      PD_THROW("The input data pointer is null.");
    }
  }

79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98
  void AssignDataFromTensor(const T& tensor) {
    size_t n = tensor.numel();
    array_.reserve(n);
    switch (tensor.dtype()) {
      case DataType::INT32:
        AssignData(tensor.template data<int32_t>(), n);
        break;
      case DataType::INT64:
        AssignData(tensor.template data<int64_t>(), n);
        break;
      default:
        PD_THROW(
            "Data type error. Currently, The data type of IntArrayBase "
            "only supports Tensor with int32 and int64, "
            "but now received `",
            tensor.dtype(),
            "`.");
    }
  }

99 100 101 102
 private:
  // TODO(zhangyunfei) Replace std::vector with a more efficient container
  // structure.
  std::vector<int64_t> array_;
103
  bool is_from_tensor_{false};
104 105
};

106 107
using IntArray =
    paddle::experimental::IntArrayBase<paddle::experimental::Tensor>;
108 109 110 111

}  // namespace experimental
}  // namespace paddle

112
namespace phi {
113 114

class DenseTensor;
115
using IntArray = paddle::experimental::IntArrayBase<DenseTensor>;
116

117
}  // namespace phi