views.hpp

Go to the documentation of this file.

// Copyright 2025 Eric Hermosis
//
// This file is part of the Tannic Tensor Library.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
 
#ifndef VIEWS_HPP
#define VIEWS_HPP 
 
#include <utility> 
#include <algorithm>
#include <numeric>
#include <vector>
 
#include "types.hpp"
#include "traits.hpp"
#include "shape.hpp"
#include "strides.hpp"
#include "concepts.hpp"
#include "exceptions.hpp"
 
namespace tannic {
 
class Tensor;
 
namespace expression {   
   
template<Expression Source>
class View {
public:  
 
    template<Integral... Sizes>  
    constexpr View(typename Trait<Source>::Reference source, Sizes... sizes)
    :   source_(source)
    { 
        std::array<long long, sizeof...(Sizes)> requested{ static_cast<long long>(sizes)... };
 
        std::size_t source_elements = std::accumulate(
            source.shape().begin(), source.shape().end(), 1ULL, std::multiplies<>{}
        );
 
        // Step 3. Handle -1 (infer axis)
        int infer_axis = -1;
        std::size_t known_product = 1;
        for (std::size_t i = 0; i < requested.size(); ++i) {
            if (requested[i] == -1) {
                if (infer_axis != -1)
                    throw Exception("Only one dimension can be inferred (-1) in view");
                infer_axis = static_cast<int>(i);
            } else if (requested[i] < 0) {
                throw Exception("Invalid negative dimension in view");
            } else {
                known_product *= static_cast<std::size_t>(requested[i]);
            }
        }
 
        if (infer_axis != -1) {
            if (source_elements % known_product != 0)
                throw Exception("Cannot infer dimension: source elements not divisible");
            requested[infer_axis] = static_cast<long long>(source_elements / known_product);
        }
 
        for (auto v : requested) {
            shape_.expand(static_cast<std::size_t>(v));
        }
 
        std::size_t new_elements = std::accumulate(shape_.begin(), shape_.end(), 1ULL, std::multiplies<>{});
        if (new_elements != source_elements)
            throw Exception("Shape mismatch: view must preserve total number of elements");
 
        strides_ = Strides(shape_);
    }
 
 
    constexpr type dtype() const {
        return source_.dtype();
    }
 
    constexpr Shape const& shape() const {
        return shape_;
    }
 
    constexpr Strides const& strides() const {
        return strides_;
    }
 
    std::ptrdiff_t offset() const {
        return source_.offset();
    }
 
    Tensor forward() const;
 
private:
    Shape shape_;
    Strides strides_;
    typename Trait<Source>::Reference source_;                
};   
 
 
template<Expression Source>
class Transpose {
public: 
 
    constexpr Transpose(typename Trait<Source>::Reference source, std::pair<int, int> dimensions)
    :   shape_(source.shape()) 
    ,   strides_(source.strides())
    ,   source_(source)
    ,   dimensions_(dimensions) {  
        auto rank = source.shape().rank();    
        std::swap(shape_[indexing::normalize(dimensions.first, rank)], shape_[indexing::normalize(dimensions.second, rank)]); 
        std::swap(strides_[indexing::normalize(dimensions.first, rank)], strides_[indexing::normalize(dimensions.second, rank)]);   
    }
 
    constexpr type dtype() const {
        return source_.dtype();
    } 
        
    constexpr Shape const& shape() const {
        return shape_;
    } 
 
    constexpr Strides const& strides() const {
        return strides_;
    }  
 
    std::ptrdiff_t offset() const {
        return source_.offset();
    }
 
    Tensor forward() const;
 
private:
    Shape shape_; 
    Strides strides_;
    typename Trait<Source>::Reference source_; 
    std::pair<int, int> dimensions_;
};   
 
 
template<Expression Source, Integral ... Indexes>
class Permutation {
public:
    constexpr Permutation(typename Trait<Source>::Reference source, std::tuple<Indexes...> indexes) 
    :   source_(source)
    {
        if (sizeof...(Indexes) != source_.shape().rank()) {
            throw Exception("Permutation rank must equal tensor rank");
        }
 
        std::apply([&](auto... indexes) {
            (([&]{
                int dimension = indexing::normalize(indexes, source_.shape().rank());
                shape_.expand(source_.shape()[dimension]);
                strides_.expand(source_.strides()[dimension]);
            }()), ...);
        }, indexes);
    }
 
    constexpr type dtype() const { 
        return source_.dtype(); 
    }
 
    constexpr Shape const& shape() const { 
        return shape_; 
    }
 
    constexpr Strides const& strides() const { 
        return strides_; 
    }
 
    std::ptrdiff_t offset() const { 
        return source_.offset(); 
    } 
 
    Tensor forward() const;
 
private:
    Shape shape_{};
    Strides strides_{};
    typename Trait<Source>::Reference source_;              
};
 
 
 
template<Expression Source>
class Expansion {
public:
    template<Integral... Sizes>
    constexpr Expansion(typename Trait<Source>::Reference source, Sizes... sizes)
    :   source_(source) {
        std::array<long long, sizeof...(Sizes)> requested{ static_cast<long long>(sizes)... }; 
 
        if (requested.size() < source.shape().rank()) 
            throw Exception("Expansion target rank must be >= source rank");
 
        std::size_t offset = requested.size() - source.shape().rank();
        for (std::size_t dimension = 0; dimension < requested.size(); ++dimension) {
            long long index = requested[dimension];
            std::size_t target;
 
            if (index == -1) {
                if (dimension < offset) {
                    throw Exception("Cannot use -1 for new leading dimensions");
                }
                target = source.shape()[dimension - offset];
            } else if (index <= 0) {
                throw Exception("Expansion size must be positive or -1");
            } else {
                target = static_cast<std::size_t>(index);
            }
 
            if (dimension < offset) { 
                shape_.expand(target);
                strides_.expand(0);
            } 
            
            else { 
                if (source.shape()[dimension - offset] == 1 && target > 1) {
                    shape_.expand(target);
                    strides_.expand(0);  // broadcast
                } else if (source.shape()[dimension - offset] == target) {
                    shape_.expand(target);
                    strides_.expand(source.strides()[dimension - offset]);
                } else {
                    throw Exception("Expansion only allows -1 (keep) or broadcasting singleton dims");
                }
            }
        }
    }
 
 
    constexpr type dtype() const { 
        return source_.dtype(); 
    }
 
    constexpr Shape const& shape() const { 
        return shape_; 
    }
 
 
    constexpr Strides const& strides() const { 
        return strides_; 
    }
 
    std::ptrdiff_t offset() const { 
        return source_.offset(); 
    }
 
    Tensor forward() const;
 
private:
    Shape shape_;
    Strides strides_;
    typename Trait<Source>::Reference source_;
};
 
template<Expression Source>
class Squeeze {
public:
    constexpr Squeeze(typename Trait<Source>::Reference source)
    : source_(source) {
        for (auto dimension = 0; dimension < source.shape().rank(); ++dimension) {
            if (source.shape()[dimension] != 1) {
                shape_.expand(source.shape()[dimension]);
                strides_.expand(source.strides()[dimension]);
            }
        }
    }
 
    constexpr type dtype() const { 
        return source_.dtype(); 
    }
 
    constexpr Shape const& shape() const { 
        return shape_; 
    }
 
    constexpr Strides const& strides() const { 
        return strides_; 
    }
 
    std::ptrdiff_t offset() const { 
        return source_.offset(); 
    }
 
    Tensor forward() const;
 
private:
    Shape shape_{};
    Strides strides_{};
    typename Trait<Source>::Reference source_;
};
 
 
 
template<Expression Source>
class Unsqueeze {
public:
    template<Integral... Axes>
    constexpr Unsqueeze(typename Trait<Source>::Reference source, Axes... axes)
    : source_(source) {
        auto rank = source.shape().rank();
        std::vector<std::size_t> normalized{ static_cast<std::size_t>(indexing::normalize(axes, rank + sizeof...(axes)))... };
        std::sort(normalized.begin(), normalized.end());
 
        size_t dimensions = rank + normalized.size();
        size_t index = 0;
        size_t axis = 0;
 
        for (auto dimension = 0; dimension < dimensions; ++dimension) {
            if (axis < normalized.size() && dimension == normalized[axis]) {
                shape_.expand(1);
                strides_.expand( (index < source.strides().rank()) ? source.strides()[index] : 1 );
                ++axis;
            } else {
                shape_.expand(source.shape()[index]);
                strides_.expand(source.strides()[index]);
                ++index;
            }
        }
    }
 
    constexpr type dtype() const { 
        return source_.dtype(); 
    }
 
    constexpr Shape const& shape() const { 
        return shape_; 
    }
 
    constexpr Strides const& strides() const { 
        return strides_; 
    }
 
    std::ptrdiff_t offset() const { 
        return source_.offset(); 
    }
 
    Tensor forward() const;
 
private:
    Shape shape_{};
    Strides strides_{};
    typename Trait<Source>::Reference source_;
};
 
  
template<Expression Source>
class Flatten {
public:
    constexpr Flatten(typename Trait<Source>::Reference source, int start = 0, int end = -1)
    :   source_(source) {
        auto rank = source.shape().rank();
 
        start = indexing::normalize(start, rank);
        end   = indexing::normalize(end, rank);
 
        if (start > end) {
            throw Exception("Flatten requires start_dim <= end_dim");
        }
 
        for (int dimension = 0; dimension < start; ++dimension) {
            shape_.expand(source.shape()[dimension]);
            strides_.expand(source.strides()[dimension]);
        }
 
        std::size_t flattened = 1;
        for (int dimension = start; dimension <= end; ++dimension) {
            flattened *= source.shape()[dimension];
        }
        shape_.expand(flattened);
        strides_.expand(source.strides()[end]);  
 
        for (int dimension = end + 1; dimension < rank; ++dimension) {
            shape_.expand(source.shape()[dimension]);
            strides_.expand(source.strides()[dimension]);
        }
    }
 
    constexpr type dtype() const { return source_.dtype(); }
 
    constexpr Shape const& shape() const { 
        return shape_; 
    }
 
    constexpr Strides const& strides() const { 
        return strides_; 
    }
 
    std::ptrdiff_t offset() const { return source_.offset(); }
 
    Tensor forward() const;
 
private:
    Shape shape_{};
    Strides strides_{};
    typename Trait<Source>::Reference source_;
};
 
 
 
  
/*
----------------------------------------------------------------------------------------------------
*/
 
 
template<Expression Source, Integral ... Indexes>
constexpr auto view(Source&& source, Indexes ... indexes) {
    return View<Source>(
        std::forward<Source>(source), indexes...
    );
} 
 
 
template<Expression Source>
constexpr auto transpose(Source&& source, int first, int second) {
    return Transpose<Source>(
        std::forward<Source>(source),
        std::make_pair(first, second)
    );
} 
 
template<Expression Source, Integral ... Indexes>
constexpr auto permute(Source&& source, Indexes... indexes) {
    return Permutation<Source, Indexes...>(
        std::forward<Source>(source), 
        std::make_tuple(indexes...)
    );
}
 
  
template<Expression Source, Integral... Sizes>
constexpr auto expand(Source&& source, Sizes... sizes) {
    return Expansion<Source>(std::forward<Source>(source), sizes...);
}
 
 
template<Expression Source>
constexpr auto squeeze(Source&& source) {
    return Squeeze<Source>(std::forward<Source>(source));
}
 
 
template<Expression Source, Integral... Axes>
constexpr auto unsqueeze(Source&& source, Axes... axes) {
    return Unsqueeze<Source>(std::forward<Source>(source), axes...);
}
 
 
template<Expression Source>
constexpr auto flatten(Source&& source, int start = 0, int end   = -1) {
    return Flatten<Source>(std::forward<Source>(source), start, end);
}
 
 
} // namespace expression
 
using expression::view;
using expression::transpose;
using expression::permute;
using expression::expand;
using expression::squeeze;
using expression::unsqueeze;
using expression::flatten;
 
} // namespace tannic
 
#endif // VIEWS_HPP