--- /dev/null
+#ifndef _NNC_CORE_LINALG_SHAPE_RANGE_H_
+#define _NNC_CORE_LINALG_SHAPE_RANGE_H_
+
+#include "nncc/core/ADT/tensor/Shape.h"
+#include "nncc/core/ADT/tensor/Index.h"
+
+namespace nncc
+{
+namespace contrib
+{
+namespace core
+{
+namespace data
+{
+
+using nncc::core::ADT::tensor::Shape;
+using nncc::core::ADT::tensor::Index;
+
+class ShapeIter :
+ public std::iterator<std::forward_iterator_tag, Index, size_t, Index*, Index&> {
+ public:
+ ShapeIter& operator++() {
+ auto* pidx = &_index.at(0);
+ auto* pshape = &_shape.dim(0);
+ auto rank = _shape.rank();
+ int c = rank - 1;
+ pidx[c]++;
+ while( (pidx[c] > pshape[c] - 1) && (c > 0) ) {
+ pidx[c] = 0;
+ pidx[--c]++;
+ }
+ _pos++;
+ return *this;
+ }
+
+ ShapeIter operator++(int) {
+ ShapeIter it = *this;
+ ++*this;
+ return it;
+ }
+
+ const Index& operator*() const {
+ return _index;
+ }
+
+ bool operator!=(ShapeIter& iter) {
+ assert(iter._index.rank() == _index.rank());
+ assert(iter._shape == _shape);
+ return _pos != iter._pos;
+ }
+
+ private:
+ explicit ShapeIter(Shape &_shape, uint32_t pos) : _pos(pos), _shape(_shape) {
+ _index.resize(_shape.rank());
+ _index.fill(0);
+ }
+
+ uint32_t _pos;
+ Index _index;
+ Shape& _shape;
+
+ friend class ShapeRange;
+};
+
+class ShapeRange {
+ public:
+ explicit ShapeRange(const Shape &shape) : _shape(const_cast<Shape&>(shape))
+ {}
+
+ ShapeIter begin() {
+ return ShapeIter(_shape, 0);
+ }
+
+ ShapeIter end() {
+ uint32_t _end = 1;
+ for( uint32_t d = 0; d < _shape.rank(); ++d ) _end *= _shape.dim(d);
+ return ShapeIter(_shape, _end);
+ }
+
+ bool contains(const Index& idx) {
+ assert(idx.rank() == _shape.rank());
+ bool res = true;
+ for(uint32_t d = 0; d < idx.rank(); ++d ) {
+ res &= idx.at(d) < _shape.dim(d);
+ }
+ return res;
+ }
+
+ private:
+ Shape& _shape;
+};
+
+} // namespace data
+} // namespace core
+} // namespace contrib
+} // namespace nncc
+
+#endif //_NNC_CORE_LINALG_SHAPE_RANGE_H_
--- /dev/null
+#include "gtest/gtest.h"
+#include "nnc/core/linalg/ShapeRange.h"
+
+namespace {
+
+using namespace nncc::contrib::core::data;
+
+struct ParamType {
+ uint32_t actual_length;
+ Shape shape;
+
+ template<typename ...Args>
+ explicit ParamType(uint32_t actual_len, Args&& ...args) : actual_length(actual_len), shape({static_cast<uint32_t>(args)...}) {}
+};
+
+class ShapeIteratorTest : public ::testing::TestWithParam<ParamType> {
+
+};
+
+TEST_P(ShapeIteratorTest, ElementCount) {
+ Shape sh(GetParam().shape);
+ ShapeRange r(sh);
+
+ uint32_t cnt = 0;
+ for( auto& idx : r ) {
+ (void)idx;
+ cnt++;
+ }
+
+ ASSERT_EQ(cnt, GetParam().actual_length);
+}
+
+
+std::vector<ParamType> test_data {
+ ParamType{6, 1,2,3},
+ ParamType{16, 2,2,4},
+ ParamType{1, 1,1,1,1,1},
+ ParamType{5, 5,1,1,1,1,1}
+};
+
+
+INSTANTIATE_TEST_CASE_P(SimpleInput, ShapeIteratorTest, ::testing::ValuesIn(test_data));
+
+}