VaryingShape Class — pytorch Architecture

Source Code

aten/src/ATen/core/jit_type.h lines 494–565

template <typename T>
struct VaryingShape {
  using ListOfOptionalElements = std::vector<std::optional<T>>;
  VaryingShape(const std::vector<T>& vec)
      : VaryingShape(ListOfOptionalElements(vec.begin(), vec.end())) {}

  VaryingShape(c10::ArrayRef<T> vec)
      : VaryingShape(ListOfOptionalElements(vec.begin(), vec.end())) {}

  VaryingShape(std::optional<size_t> size = std::nullopt) : dims_(std::nullopt) {
    if (size) {
      dims_ = ListOfOptionalElements(*size);
    }
  }

  VaryingShape(ListOfOptionalElements dims) : dims_(std::move(dims)) {}

  VaryingShape(size_t size) : VaryingShape(std::optional<size_t>(size)) {}

  bool operator==(const VaryingShape& other) const {
    return dims_ == other.dims_;
  }

  const std::optional<T> &operator[](size_t i) const {
    TORCH_CHECK(dims_, "Rank isn't fixed");
    return (*dims_).at(i);
  }

  std::optional<size_t> size() const {
    if (!dims_) {
      return std::nullopt;
    }
    const auto& dims = dims_.value();
    return dims.size();
  }

  const std::optional<ListOfOptionalElements>& sizes() const {
    return dims_;
  }

  TORCH_API VaryingShape merge(const VaryingShape& other) const;

  std::optional<std::vector<T>> concrete_sizes() const {
    if (!dims_) {
      return std::nullopt;
    }
    std::vector<T> sizes;
    sizes.reserve(dims_.value().size());
    for (auto d : *dims_) {
      if (!d) {
        return std::nullopt;
      }
      sizes.push_back(d.value());
    }
    return sizes;
  }

  bool isComplete() const {
    if (!dims_) {
      return false;
    }
    for (auto d : *dims_) {
      if (!d || !detail::isComplete(*d)) {
        return false;
      }
    }
    return true;
  }

 private:
  std::optional<ListOfOptionalElements> dims_;
};