shape.hpp

Go to the documentation of this file.

 
#pragma once
 
#include <cstdint>
#include <initializer_list>
#include <limits>
#include <numeric>
#include <stdexcept>
#include <vector>
 
#include "pulsatrix/assert.hpp"
 
namespace pulsatrix {
 
class Shape {
public:
    Shape(std::initializer_list<int64_t> dims) : dims_(dims) {
        for (int64_t d : dims_) {
            if (d < 0) {
                throw std::invalid_argument("Shape: dimensions must be non-negative");
            }
        }
    }
 
    explicit Shape(const std::vector<int64_t>& dims) : dims_(dims) {
        for (int64_t d : dims_) {
            if (d < 0) {
                throw std::invalid_argument("Shape: dimensions must be non-negative");
            }
        }
    }
 
    [[nodiscard]] int64_t rank() const { return static_cast<int64_t>(dims_.size()); }
 
    [[nodiscard]] int64_t numel() const {
        int64_t result = 1;
        for (int64_t d : dims_) {
            if (d != 0 && result > std::numeric_limits<int64_t>::max() / (d == 0 ? 1 : d)) {
                throw std::overflow_error("Shape::numel: dimension product overflows int64_t");
            }
            result *= d;
        }
        return result;
    }
 
    [[nodiscard]] int64_t dim(size_t index) const {
        PULSATRIX_ASSERT(index < dims_.size());
        return dims_[index];
    }
 
    [[nodiscard]] bool is_reshape_compatible(const Shape& other) const { return numel() == other.numel(); }
 
    [[nodiscard]] bool operator==(const Shape& other) const { return dims_ == other.dims_; }
    [[nodiscard]] bool operator!=(const Shape& other) const { return !(*this == other); }
 
private:
    std::vector<int64_t> dims_;
};
 
}  // namespace pulsatrix