tensor_shape.h

Go to the documentation of this file.

#ifndef ATEN_CORE_TENSOR_SHAPE_H_
#define ATEN_CORE_TENSOR_SHAPE_H_
 
#include <vector>
#include <iostream>
#include <initializer_list>
 
namespace container {
 
class TensorShape {
public:
    TensorShape();
 
    TensorShape(std::initializer_list<int64_t> dims);
 
    TensorShape(const std::vector<int64_t>& dims);
 
    TensorShape(const TensorShape& other);
 
    int64_t dim_size(int dim) const;
 
    const std::vector<int64_t>& dims() const;
 
    const std::vector<int64_t>& strides() const;
 
    unsigned int ndim() const;
 
    void set_dim_size(int dim, int64_t size);
 
    void add_dim(int64_t size);
 
    void remove_dim(int dim);
 
    int64_t NumElements() const;
 
    bool operator==(const TensorShape& other) const;
 
    bool operator!=(const TensorShape& other) const;
 
private:
    std::vector<int64_t> dims_ = {};     // Save dimension sizes of the tensor
    // Note: strides are not always equals to the dimension sizes.
    // The strides specifies the number of elements to step in each dimension when traversing a tensor.
    // There could be some sparse region in the tensor, and the strides will be larger than the dimension sizes.
    // For example, given a 2D tensor with shape [3, 4], 
    // and the strides could be [6, 1] if the actual data is stored in a 1D array with size 18 [3, 6].
    // The strides could also be [12, 3] if the actual data is stored in a 1D array with size 36 [3, 12].
    // strides can only be modified by the TensorMap object.
    std::vector<int64_t> strides_ = {};  // Save dimension strides of the tensor
 
    std::vector<int64_t> get_strides_(const std::vector<int64_t>& dim);
};
 
std::ostream& operator<<(std::ostream& os, const TensorShape& shape);
 
} // container
 
#endif  // ATEN_CORE_TENSOR_SHAPE_H_