Class ShapeIndex

tensorflow/compiler/xla/shape_util.h:67–116

An index for specifying a particular nested subshape within a shape. Used in ShapeUtil::GetSubshape and other interfaces. Shapes are recursive data structures (trees) and ShapeIndex defines a path through the tree where each element of ShapeIndex indexes into a tuple (or nested tuple) within the shape. For a non-nested tuple, an index has a single element. For example, given a 3-element tuple (a,

// ShapeIndex is a trivial wrapper around std::vector with a minimum number of
// methods implemented.
class ShapeIndex {
 public:
  ShapeIndex() = default;
  ShapeIndex(std::initializer_list<int64> init) : indices_(init) {}
  template <typename InputIt>
  ShapeIndex(InputIt start, InputIt end) : indices_(start, end) {}

  explicit ShapeIndex(ShapeIndexView v);

  bool empty() const { return indices_.empty(); }
  size_t size() const { return indices_.size(); }
  void push_back(int64 value) { indices_.push_back(value); }
  void pop_back() { indices_.pop_back(); }

  // push_front is O(n), but shapes don't usually have a ton of dimensions.
  void push_front(int64 value) { indices_.insert(indices_.begin(), value); }

  using container_type = absl::InlinedVector<int64, 2>;

  container_type::const_iterator begin() const { return indices_.begin(); }
  container_type::const_iterator end() const { return indices_.end(); }
  container_type::iterator begin() { return indices_.begin(); }
  container_type::iterator end() { return indices_.end(); }

  const int64* data() const { return indices_.data(); }

  int64 back() const { return indices_.back(); }
  int64& back() { return indices_.back(); }

  const int64& operator[](size_t i) const { return indices_[i]; }
  int64& operator[](size_t i) { return indices_[i]; }

  bool operator==(const ShapeIndex& other) const {
    return indices_ == other.indices_;
  }
  bool operator!=(const ShapeIndex& other) const { return !(*this == other); }
  bool operator<(const ShapeIndex& other) const {
    return indices_ < other.indices_;
  }

  string ToString() const;

  template <typename H>
  friend H AbslHashValue(H h, const ShapeIndex& index) {
    return H::combine(std::move(h), index.indices_);
  }

 private:
  container_type indices_;
};

// A view into a ShapeIndex as above, with the cheap/easy ability to consume the
// value at the front of the view.