reductions.hpp

Go to the documentation of this file.

// Copyright 2025 Eric Hermosis
//
// 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.
//
 
// Copyright 2025 Eric Hermosis
//
// 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 REDUCTIONS_HPP
#define REDUCTIONS_HPP 
 
#include <array> 
#include <cassert>
 
#include "concepts.hpp"
#include "types.hpp"
#include "traits.hpp"
#include "shape.hpp" 
#include "tensor.hpp" 
#include "indexing.hpp"
#include "exceptions.hpp"
 
 
namespace tannic::expression {
 
 
template<class Reducer, Expression Operand>
class Reduction {
public:
    Reducer reducer;
    typename Trait<Operand>::Reference operand;
 
    constexpr Reduction(Reducer reducer, typename Trait<Operand>::Reference operand)
    :   reducer(reducer)
    ,   operand(operand)
    ,   dtype_(reducer.reduce(operand.dtype()))
    ,   shape_(reducer.reduce(operand.shape()))
    ,   strides_(shape_) {}
 
    constexpr type dtype() const {
        return dtype_;
    }
 
    constexpr Shape const& shape() const {
        return shape_;
    } 
 
    constexpr Strides const& strides() const {
        return strides_;
    }
 
    std::ptrdiff_t offset() const {
        return 0;
    }
 
    Tensor forward() const {   
        Tensor source = operand.forward();
        Tensor result(dtype(), shape(), strides(), offset());          
        reducer.forward(source, result);
        return result;
    }
 
private:
    type dtype_;
    Shape shape_;
    Strides strides_;
};
 
 
struct Argmax {   
    int axis; 
    bool keepdim;   
 
    constexpr type reduce(type dtype) const {
        return int64;
    } 
 
    constexpr Shape reduce(Shape const& shape) const {
        if (shape.rank() == 0) 
            throw Exception("Cannot reduce scalar tensors");
 
        Shape out;
        for (size_t dim = 0; dim < shape.rank(); ++dim) {
            if (dim != static_cast<size_t>(axis)) out.expand(shape[dim]);
            else if (keepdim) out.expand(1);
        }
        return out;
    }
  
    void forward(Tensor const& input, Tensor& output) const;
}; 
 
struct Argmin {   
    int axis;    
    bool keepdim;  
 
    constexpr type reduce(type dtype) const {
        return int64;
    }
 
    constexpr Shape reduce(Shape const& shape) const {
        if (shape.rank() == 0) 
            throw Exception("Cannot reduce scalar tensors");
 
        Shape reduced; 
        for (uint8_t dimension = 0; dimension < shape.rank(); ++dimension) {
            if (dimension != static_cast<uint8_t>(axis))
                reduced.expand(shape[dimension]);
        }
        return reduced;
    } 
 
    void forward(Tensor const&, Tensor&) const; 
};
 
 
struct Argsum {
    int axis;       
    bool keepdim;  
    
    constexpr type reduce(type dtype) const {
        return dtype;   
    }
 
    constexpr Shape reduce(Shape const& shape) const {
        Shape reduced;
        for (size_t dim = 0; dim < shape.rank(); ++dim) {
            if (dim != static_cast<size_t>(axis)) {
                reduced.expand(shape[dim]);
            } else if (keepdim) {
                reduced.expand(1);  // Keep reduced dim as size 1
            }
        }
        return reduced;
    } 
 
    void forward(Tensor const& input, Tensor& output) const;
};
 
 
struct Argmean {
    int axis;
    bool keepdim = false;
    
    constexpr type reduce(type dtype) const { 
        assert(dtype == float32 | dtype == float64 && "Integral dtypes not supported.");
        return dtype;
    }
 
    constexpr Shape reduce(Shape const& shape) const {
        Shape reduced;
        for (size_t dim = 0; dim < shape.rank(); ++dim) {
            if (dim != static_cast<size_t>(axis)) {
                reduced.expand(shape[dim]);
            } else if (keepdim) {
                reduced.expand(1);
            }
        }
        return reduced;
    } 
 
    void forward(Tensor const& input, Tensor& output) const;
};
 
 
template<Expression Source>
constexpr auto argmax(Source&& source, int axis = -1, bool keepdim = false) {  
    return Reduction<Argmax, Source>{
        {indexing::normalize(axis, source.shape().rank()), keepdim}, std::forward<Source>(source) 
    };
}
 
 
template<Expression Source>
constexpr auto argmin(Source&& source, int axis = -1, bool keepdim = false) {  
    return Reduction<Argmin, Source>{ 
        {indexing::normalize(axis, source.shape().rank()), keepdim}, std::forward<Source>(source) 
    };
} 
 
template<Expression Source>
constexpr auto sum(Source&& source, int axis = -1, bool keepdim = false) {
    return Reduction<Argsum, Source>{
        {indexing::normalize(axis, source.shape().rank()), keepdim},
        std::forward<Source>(source)
    };
}
 
template<Expression Source>
constexpr auto mean(Source&& source, int axis = -1, bool keepdim = false) {
    return Reduction<Argmean, Source>{
        {indexing::normalize(axis, source.shape().rank()), keepdim},
        std::forward<Source>(source)
    };
}
 
} namespace tannic {
 
using expression::argmax;
using expression::argmin;
using expression::sum;
using expression::mean;
 
} // namespace tannic
 
#endif // REDUCTIONS_HPP