combine_scaled_embeddings.hpp

Go to the documentation of this file.

#ifndef MUMOSA_COMBINE_SCALED_EMBEDDINGS_HPP
#define MUMOSA_COMBINE_SCALED_EMBEDDINGS_HPP
 
#include <vector>
#include <cstddef>
#include <algorithm>
#include <stdexcept>
#include <cmath>
 
#include "sanisizer/sanisizer.hpp"
 
#include "utils.hpp"
 
namespace mumosa {
 
template<typename Index_, typename Input_, typename Scale_, typename Output_>
void combine_scaled_embeddings(
    const std::vector<std::size_t>& num_dims,
    const Index_ num_cells,
    const std::vector<Input_*>& embeddings,
    const std::vector<Scale_>& scaling,
    Output_* const output
) {
    const auto nembed = num_dims.size();
    if (embeddings.size() != nembed || scaling.size() != nembed) {
        throw std::runtime_error("'num_dims', 'embeddings' and 'scale' should have the same length");
    }
 
    const std::size_t ntotal = std::accumulate(num_dims.begin(), num_dims.end(), static_cast<std::size_t>(0));
    std::size_t starting_dim = 0;
 
    for (I<decltype(nembed)> e = 0; e < nembed; ++e) {
        const auto curdim = num_dims[e];
        const auto inptr = embeddings[e];
        const auto s = scaling[e];
 
        if (std::isinf(s)) {
            // If the scaling factor is infinite, it implies that the current
            // embedding is all-zero, so we just fill with zeros, and move on.
            for (Index_ c = 0; c < num_cells; ++c) {
                const auto out_offset = sanisizer::nd_offset<std::size_t>(starting_dim, ntotal, c);
                std::fill_n(output + out_offset, curdim, 0);
            }
        } else {
            for (Index_ c = 0; c < num_cells; ++c) {
                for (I<decltype(curdim)> d = 0; d < curdim; ++d) {
                    const auto out_offset = sanisizer::nd_offset<std::size_t>(starting_dim + d, ntotal, c);
                    const auto in_offset = sanisizer::nd_offset<std::size_t>(d, curdim, c);
                    output[out_offset] = inptr[in_offset] * s;
                }
            }
        }
 
        starting_dim += curdim;
    }
}
 
}
 
#endif