hyperverse 0.0.1__py3-none-any.whl

hyperdrone/_vendor/rl-tools/include/rl_tools/nn/layers/embedding/layer.h

@@ -0,0 +1,119 @@
+#include "../../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_NN_LAYERS_EMBEDDING_LAYER_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_NN_LAYERS_EMBEDDING_LAYER_H
+#include "../../../nn/activation_functions.h"
+#include "../../../utils/generic/typing.h"
+
+//#include "../../../nn/nn.h"
+#include "../../../nn/capability/capability.h"
+#include "../../../nn/parameters/parameters.h"
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools::nn::layers::embedding {
+//    template <typename LAYER_SPEC, typename INPUT_SPEC, typename OUTPUT_SPEC>
+//    constexpr bool check_input_output_f(){
+//        static_assert(length(typename INPUT_SPEC::SHAPE{}) == 1);
+//        static_assert(length(typename OUTPUT_SPEC::SHAPE{}) == 2);
+//        static_assert(get<0>(typename INPUT_SPEC::SHAPE{}) == get<0>(typename OUTPUT_SPEC::SHAPE{}));
+//        static_assert(get<1>(typename OUTPUT_SPEC::SHAPE{}) == LAYER_SPEC::OUTPUT_DIM);
+//        return true;
+//    }
+//    template <typename LAYER_SPEC, typename INPUT_SPEC, typename OUTPUT_SPEC>
+//    constexpr bool check_input_output = check_input_output_f<LAYER_SPEC, INPUT_SPEC, OUTPUT_SPEC>();
+    template <typename T_TYPE_POLICY, typename T_TI>
+    struct StandardNormalSpecification{
+        using TYPE_POLICY = T_TYPE_POLICY;
+        using T = typename TYPE_POLICY::DEFAULT;
+        using TI = T_TI;
+        static constexpr T SCALE = 1;
+    };
+    template<typename SPEC>
+    struct StandardNormal {
+    };
+    template<typename T_T, typename T_TI>
+    using DefaultInitializer = StandardNormal<StandardNormalSpecification<T_T, T_TI>>;
+
+    template <typename T_TI>
+    struct DefaultInputShapeFactory{
+        template <T_TI BATCH_SIZE>
+        using SHAPE = tensor::Shape<T_TI, BATCH_SIZE>;
+    };
+    template<typename T_TYPE_POLICY, typename T_TI, T_TI T_NUM_CLASSES, T_TI T_EMBEDDING_DIM, template <T_TI> typename T_INPUT_SHAPE = DefaultInputShapeFactory<T_TI>::template SHAPE, typename T_INITIALIZER = DefaultInitializer<T_TYPE_POLICY, T_TI>, typename T_PARAMETER_GROUP=parameters::groups::Input>
+    struct Configuration {
+        using TYPE_POLICY = T_TYPE_POLICY;
+        using TI = T_TI;
+        static constexpr TI NUM_CLASSES = T_NUM_CLASSES;
+        static constexpr TI EMBEDDING_DIM = T_EMBEDDING_DIM;
+        using INITIALIZER = T_INITIALIZER;
+        using PARAMETER_GROUP = T_PARAMETER_GROUP;
+        // Summary
+        static constexpr TI NUM_WEIGHTS = EMBEDDING_DIM * NUM_CLASSES;
+    };
+
+    template <typename T_CONFIG, typename T_CAPABILITY, typename T_INPUT_SHAPE>
+    struct Specification: T_CAPABILITY, T_CONFIG{
+        using CONFIG = T_CONFIG;
+        using CAPABILITY = T_CAPABILITY;
+        using INPUT_SHAPE = T_INPUT_SHAPE;
+        using TYPE_POLICY = typename CONFIG::TYPE_POLICY;
+        using TI = typename CONFIG::TI;
+        static_assert(length(INPUT_SHAPE{}) == 3, "The input shape of the Embedding layer must be 3 dimensional for now (sequence x batch x 1 (integer ids))");
+        static constexpr TI INPUT_DIM = get_last(INPUT_SHAPE{});
+        static constexpr TI SEQUENCE_LENGTH = get<length(INPUT_SHAPE{})-3>(INPUT_SHAPE{});
+        template <typename NEW_INPUT_SHAPE>
+        using OUTPUT_SHAPE_FACTORY = tensor::Replace<NEW_INPUT_SHAPE, CONFIG::EMBEDDING_DIM, length(NEW_INPUT_SHAPE{})-1>;
+        using OUTPUT_SHAPE = OUTPUT_SHAPE_FACTORY<INPUT_SHAPE>;
+        static constexpr TI INTERNAL_BATCH_SIZE = get<1>(INPUT_SHAPE{}); // Since the Dense layer is based on Matrices (2D Tensors) the dense layer operation is broadcasted over the leading dimensions. Hence, the actual batch size is the product of all leading dimensions, excluding the last one (containing the features). Since rl_tools::matrix_view is used for zero-cost conversion the INTERNAL_BATCH_SIZE accounts for all leading dimensions.
+    };
+
+    struct State{};
+    struct Buffer{};
+
+    template<typename T_SPEC>
+    struct LayerForward {
+        using SPEC = T_SPEC;
+        using TYPE_POLICY = typename SPEC::TYPE_POLICY;
+        using TI = typename SPEC::TI;
+        static constexpr TI NUM_CLASSES = SPEC::NUM_CLASSES;
+        static constexpr TI OUTPUT_DIM = SPEC::EMBEDDING_DIM;
+        static constexpr TI NUM_WEIGHTS = SPEC::NUM_WEIGHTS;
+        using INPUT_SHAPE = typename SPEC::INPUT_SHAPE;
+        using OUTPUT_SHAPE = typename SPEC::OUTPUT_SHAPE;
+        template <typename NEW_INPUT_SHAPE>
+        using OUTPUT_SHAPE_FACTORY = typename SPEC::template OUTPUT_SHAPE_FACTORY<NEW_INPUT_SHAPE>;
+        using WEIGHTS_SHAPE = tensor::Shape<TI, NUM_CLASSES, OUTPUT_DIM>;
+        using WEIGHTS_PARAMETER_SPEC = typename SPEC::PARAMETER_TYPE::template Specification<TYPE_POLICY, TI, WEIGHTS_SHAPE, typename SPEC::PARAMETER_GROUP, nn::parameters::categories::Weights, SPEC::DYNAMIC_ALLOCATION>;
+        typename SPEC::PARAMETER_TYPE::template Instance<WEIGHTS_PARAMETER_SPEC> weights;
+
+        template<bool DYNAMIC_ALLOCATION=true>
+        using State = embedding::State;
+        template<bool DYNAMIC_ALLOCATION=true>
+        using Buffer = embedding::Buffer;
+    };
+    template<typename SPEC>
+    struct LayerBackward: public LayerForward<SPEC>{
+    };
+    template<typename SPEC>
+    struct LayerGradient: public LayerBackward<SPEC>{
+        // This layer supports backpropagation wrt its input but including its weights (for this it stores the intermediate outputs in addition to the pre_activations because they determine the gradient wrt the weights of the following layer)
+        using T_ACC = typename SPEC::TYPE_POLICY::template GET<numeric_types::categories::Accumulator>;
+        using OUTPUT_CONTAINER_SPEC = tensor::Specification<T_ACC, typename SPEC::TI, typename SPEC::OUTPUT_SHAPE, SPEC::DYNAMIC_ALLOCATION>;
+        using OUTPUT_CONTAINER_TYPE = Tensor<OUTPUT_CONTAINER_SPEC>;
+        OUTPUT_CONTAINER_TYPE output;
+    };
+    template<typename CONFIG, typename CAPABILITY, typename INPUT_SHAPE>
+    using Layer =
+        typename utils::typing::conditional_t<CAPABILITY::TAG == nn::LayerCapability::Forward, LayerForward<Specification<CONFIG, CAPABILITY, INPUT_SHAPE>>,
+        typename utils::typing::conditional_t<CAPABILITY::TAG == nn::LayerCapability::Backward, LayerBackward<Specification<CONFIG, CAPABILITY, INPUT_SHAPE>>,
+        typename utils::typing::conditional_t<CAPABILITY::TAG == nn::LayerCapability::Gradient, LayerGradient<Specification<CONFIG, CAPABILITY, INPUT_SHAPE>>, void>>>;
+
+    template <typename CONFIG>
+    struct BindConfiguration{
+        template <typename CAPABILITY, typename INPUT_SHAPE>
+        using Layer = nn::layers::embedding::Layer<CONFIG, CAPABILITY, INPUT_SHAPE>;
+    };
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+
+#endif

hyperdrone/_vendor/rl-tools/include/rl_tools/nn/layers/embedding/operations_generic.h

@@ -0,0 +1,223 @@
+#include "../../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_NN_LAYERS_EMBEDDING_OPERATIONS_GENERIC_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_NN_LAYERS_EMBEDDING_OPERATIONS_GENERIC_H
+
+#include "../../../containers/matrix/matrix.h"
+#include "../../../nn/parameters/operations_generic.h"
+
+#include "layer.h"
+#include "../../../rl_tools.h"
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools{
+    template<typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void malloc(DEVICE& device, nn::layers::embedding::LayerForward<SPEC>& layer) {
+        malloc(device, layer.weights);
+    }
+    template<typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void free(DEVICE& device, nn::layers::embedding::LayerForward<SPEC>& layer) {
+        free(device, layer.weights);
+    }
+    template<typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void malloc(DEVICE& device, nn::layers::embedding::LayerBackward<SPEC>& layer) {
+        malloc(device, (nn::layers::embedding::LayerForward<SPEC>&) layer);
+    }
+    template<typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void free(DEVICE& device, nn::layers::embedding::LayerBackward<SPEC>& layer) {
+        free(device, (nn::layers::embedding::LayerForward<SPEC>&) layer);
+    }
+    template<typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void malloc(DEVICE& device, nn::layers::embedding::LayerGradient<SPEC>& layer) {
+        malloc(device, (nn::layers::embedding::LayerBackward<SPEC>&) layer);
+        malloc(device, layer.output);
+    }
+    template<typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void free(DEVICE& device, nn::layers::embedding::LayerGradient<SPEC>& layer) {
+        free(device, (nn::layers::embedding::LayerBackward<SPEC>&) layer);
+        free(device, layer.output);
+    }
+    template<typename DEVICE>
+    RL_TOOLS_FUNCTION_PLACEMENT void malloc(DEVICE& device, nn::layers::embedding::State& state) { } // no-op
+    template <typename SOURCE_DEVICE, typename TARGET_DEVICE>
+    RL_TOOLS_FUNCTION_PLACEMENT void copy(SOURCE_DEVICE& source_device, TARGET_DEVICE& target_device, nn::layers::embedding::State& source, nn::layers::embedding::State& target){}
+    template<typename DEVICE, typename SPEC, typename RNG, typename MODE>
+    RL_TOOLS_FUNCTION_PLACEMENT void reset(DEVICE& device, const nn::layers::embedding::LayerForward<SPEC>& layer, nn::layers::embedding::State& state, RNG&, Mode<MODE> mode = Mode<mode::Default<>>{}) { } // no-op
+    template<typename DEVICE>
+    RL_TOOLS_FUNCTION_PLACEMENT void free(DEVICE& device, nn::layers::embedding::State& state) { } // no-op
+    template<typename DEVICE>
+    RL_TOOLS_FUNCTION_PLACEMENT void malloc(DEVICE& device, nn::layers::embedding::Buffer& buffer) { } // no-op
+    template<typename DEVICE>
+    RL_TOOLS_FUNCTION_PLACEMENT void free(DEVICE& device, nn::layers::embedding::Buffer& buffer) { } // no-op
+
+    template<typename DEVICE, typename SPEC, typename INITIALIZER_SPEC, typename RNG>
+    RL_TOOLS_FUNCTION_PLACEMENT void init_weights(DEVICE& device, nn::layers::embedding::LayerForward<SPEC>& layer, const nn::layers::embedding::StandardNormal<INITIALIZER_SPEC>& initializer, RNG& rng) {
+        using T = typename decltype(layer.weights.parameters)::SPEC::T;
+        using TI = typename SPEC::TI;
+        for(TI class_i = 0; class_i < SPEC::NUM_CLASSES; class_i++){
+            for(TI dim_i = 0; dim_i < SPEC::EMBEDDING_DIM; dim_i++){
+                T value = random::normal_distribution::sample(device.random, (T)0, (T)1, rng);
+                set(device, layer.weights.parameters, value, class_i, dim_i);
+            }
+        }
+    }
+    template<typename DEVICE, typename SPEC, typename RNG>
+    RL_TOOLS_FUNCTION_PLACEMENT void init_weights(DEVICE& device, nn::layers::embedding::LayerForward<SPEC>& layer, RNG& rng) {
+        init_weights(device, layer, typename SPEC::INITIALIZER{}, rng);
+    }
+
+//#ifndef RL_TOOLS_NN_DISABLE_GENERIC_FORWARD_BACKWARD
+    template<typename DEVICE, typename LAYER_SPEC, typename INPUT_SPEC, typename OUTPUT_SPEC, typename RNG, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT void evaluate(DEVICE& device, const nn::layers::embedding::LayerForward<LAYER_SPEC>& layer, const Tensor<INPUT_SPEC>& input, Tensor<OUTPUT_SPEC>& output, nn::layers::embedding::Buffer&, RNG& rng, const Mode<MODE>& mode = Mode<mode::Default<>>{}) {
+//        static_assert(nn::layers::embedding::check_input_output<LAYER_SPEC, INPUT_SPEC, OUTPUT_SPEC>);
+        // Warning do not use the same buffer for input and output!
+        using TI = typename DEVICE::index_t;
+//        auto input_flat = reshape
+        constexpr TI INPUT_ELEMENTS = get<0>(tensor::CumulativeProduct<typename INPUT_SPEC::SHAPE>{});
+        auto input_view = reshape_row_major(device, input, tensor::Shape<TI, INPUT_ELEMENTS>{});
+        auto output_view = reshape_row_major(device, output, tensor::Shape<TI, INPUT_ELEMENTS, LAYER_SPEC::EMBEDDING_DIM>{});
+        for(TI batch_i=0; batch_i < get<0>(typename decltype(input_view)::SPEC::SHAPE{}); batch_i++){
+            auto index = get(device, input_view, batch_i);
+            auto embedding = view(device, layer.weights.parameters, index, tensor::ViewSpec<0>{});
+            auto output_row = view(device, output_view, batch_i, tensor::ViewSpec<0>{});
+            copy(device, device, embedding, output_row);
+        }
+    }
+    template<typename DEVICE, typename LAYER_SPEC, typename INPUT_SPEC, typename OUTPUT_SPEC, typename RNG, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT void evaluate_step(DEVICE& device, const nn::layers::embedding::LayerForward<LAYER_SPEC>& layer, const Tensor<INPUT_SPEC>& input, nn::layers::embedding::State& state, Tensor<OUTPUT_SPEC>& output, nn::layers::embedding::Buffer& buffer, RNG& rng, const Mode<MODE>& mode = Mode<mode::Default<>>{}) {
+        evaluate(device, layer, input, output, buffer, rng, mode);
+    }
+
+    template<typename DEVICE, typename LAYER_SPEC, typename INPUT_SPEC, typename OUTPUT_SPEC, typename RNG, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT void forward(DEVICE& device, nn::layers::embedding::LayerBackward<LAYER_SPEC>& layer, const Tensor<INPUT_SPEC>& input, Tensor<OUTPUT_SPEC>& output, nn::layers::embedding::Buffer& buffer, RNG& rng, const Mode<MODE>& mode = Mode<mode::Default<>>{}){
+        evaluate(device, static_cast<nn::layers::embedding::LayerForward<LAYER_SPEC>&>(layer), input, output, buffer, rng, mode);
+    }
+//#endif
+
+    template<typename DEVICE, typename LAYER_SPEC, typename INPUT_SPEC, typename RNG, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT void forward(DEVICE& device, nn::layers::embedding::LayerGradient<LAYER_SPEC>& layer, const Tensor<INPUT_SPEC>& input, nn::layers::embedding::Buffer& buffer, RNG& rng, const Mode<MODE>& mode = Mode<mode::Default<>>{}) {
+        forward(device, static_cast<nn::layers::embedding::LayerBackward<LAYER_SPEC>&>(layer), input, layer.output, buffer, rng, mode);
+    }
+    template<typename DEVICE, typename LAYER_SPEC, typename INPUT_SPEC, typename OUTPUT_SPEC, typename RNG, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT void forward(DEVICE& device, nn::layers::embedding::LayerGradient<LAYER_SPEC>& layer, const Tensor<INPUT_SPEC>& input, Tensor<OUTPUT_SPEC>& output, nn::layers::embedding::Buffer& buffer, RNG& rng, const Mode<MODE>& mode = Mode<mode::Default<>>{}) {
+        forward(device, layer, input, buffer, rng, mode);
+        copy(device, device, layer.output, output);
+    }
+
+//#ifndef RL_TOOLS_NN_DISABLE_GENERIC_FORWARD_BACKWARD
+    // backward_input / backward_full are not supported because the inputs are discrete classes
+    template<typename DEVICE, typename LAYER_SPEC, typename INPUT_SPEC, typename D_OUTPUT_SPEC, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT void backward(DEVICE& device, nn::layers::embedding::LayerGradient<LAYER_SPEC>& layer, const Tensor<INPUT_SPEC>& input, Tensor<D_OUTPUT_SPEC>& d_output, nn::layers::embedding::Buffer&, const Mode<MODE>& mode = Mode<mode::Default<>>{}) {
+        using T = typename decltype(layer.weights.gradient)::SPEC::T;
+        using TI = typename DEVICE::index_t;
+        constexpr TI EMBEDDING_DIM = LAYER_SPEC::EMBEDDING_DIM;
+
+        constexpr TI INPUT_ELEMENTS = get<0>(tensor::CumulativeProduct<typename INPUT_SPEC::SHAPE>{});
+        auto input_view = reshape_row_major(device, input, tensor::Shape<TI, INPUT_ELEMENTS>{});
+        auto d_output_view = reshape_row_major(device, d_output, tensor::Shape<TI, INPUT_ELEMENTS, EMBEDDING_DIM>{});
+
+        for(TI element_i=0; element_i < INPUT_ELEMENTS; element_i++){
+            for(TI output_i = 0; output_i < EMBEDDING_DIM; output_i++){
+                typename INPUT_SPEC::T class_id = get(device, input_view, element_i);
+                T d_output_value = get(device, d_output_view, element_i, output_i);
+                T gradient_value = get(device, layer.weights.gradient, (TI)class_id, output_i);
+                gradient_value += d_output_value;
+                set(device, layer.weights.gradient, gradient_value, (TI)class_id, output_i);
+            }
+        }
+    }
+//#endif
+    template<typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void zero_gradient(DEVICE& device, nn::layers::embedding::LayerGradient<SPEC>& layer) {
+        zero_gradient(device, layer.weights);
+    }
+    template<typename DEVICE, typename SOURCE_SPEC, typename TARGET_SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void add_gradient(DEVICE& device, nn::layers::embedding::LayerGradient<SOURCE_SPEC>& source, nn::layers::embedding::LayerGradient<TARGET_SPEC>& target) {
+        add_gradient(device, source.weights, target.weights);
+    }
+    template<typename SOURCE_DEVICE, typename TARGET_DEVICE, typename SOURCE_SPEC, typename TARGET_SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void copy_gradient(SOURCE_DEVICE& source_device, TARGET_DEVICE& target_device, const nn::layers::embedding::LayerGradient<SOURCE_SPEC>& source, nn::layers::embedding::LayerGradient<TARGET_SPEC>& target) {
+        copy_gradient(source_device, target_device, source.weights, target.weights);
+    }
+    template<typename DEVICE, typename SPEC, typename OPTIMIZER>
+    RL_TOOLS_FUNCTION_PLACEMENT void update(DEVICE& device, nn::layers::embedding::LayerGradient<SPEC>& layer, OPTIMIZER& optimizer){
+        update(device, layer.weights, optimizer);
+    }
+
+    template<typename DEVICE, typename SPEC, typename OPTIMIZER>
+    RL_TOOLS_FUNCTION_PLACEMENT void _reset_optimizer_state(DEVICE& device, nn::layers::embedding::LayerGradient<SPEC>& layer, OPTIMIZER& optimizer) {
+        _reset_optimizer_state(device, layer.weights, optimizer);
+    }
+
+    template<typename SOURCE_DEVICE, typename TARGET_DEVICE, typename SOURCE_SPEC, typename TARGET_SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void copy(SOURCE_DEVICE& source_device, TARGET_DEVICE& target_device, const  nn::layers::embedding::LayerForward<SOURCE_SPEC>& source, nn::layers::embedding::LayerForward<TARGET_SPEC>& target){
+        copy(source_device, target_device, source.weights, target.weights);
+    }
+    template<typename SOURCE_DEVICE, typename TARGET_DEVICE, typename SOURCE_SPEC, typename TARGET_SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void copy(SOURCE_DEVICE& source_device, TARGET_DEVICE& target_device, const nn::layers::embedding::LayerBackward<SOURCE_SPEC>& source, nn::layers::embedding::LayerBackward<TARGET_SPEC>& target){
+        copy(source_device, target_device, static_cast<const nn::layers::embedding::LayerForward<SOURCE_SPEC>&>(source), static_cast<nn::layers::embedding::LayerForward<TARGET_SPEC>&>(target));
+    }
+    template<typename SOURCE_DEVICE, typename TARGET_DEVICE, typename SOURCE_SPEC, typename TARGET_SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void copy(SOURCE_DEVICE& source_device, TARGET_DEVICE& target_device, const nn::layers::embedding::LayerGradient<SOURCE_SPEC>& source, nn::layers::embedding::LayerGradient<TARGET_SPEC>& target){
+        copy(source_device, target_device, static_cast<const nn::layers::embedding::LayerBackward<SOURCE_SPEC>&>(source), static_cast<nn::layers::embedding::LayerBackward<TARGET_SPEC>&>(target));
+    }
+    template <typename DEVICE, typename SPEC_1, typename SPEC_2>
+    RL_TOOLS_FUNCTION_PLACEMENT typename SPEC_1::TYPE_POLICY::DEFAULT abs_diff(DEVICE& device, const rl_tools::nn::layers::embedding::LayerForward<SPEC_1>& l1, const rl_tools::nn::layers::embedding::LayerForward<SPEC_2>& l2) {
+        return abs_diff(device, l1.weights, l2.weights);
+    }
+    template <typename DEVICE, typename SPEC_1, typename SPEC_2>
+    RL_TOOLS_FUNCTION_PLACEMENT typename SPEC_1::T abs_diff(DEVICE& device, const rl_tools::nn::layers::embedding::LayerBackward<SPEC_1>& l1, const rl_tools::nn::layers::embedding::LayerBackward<SPEC_2>& l2) {
+        using T = typename SPEC_1::T;
+        return abs_diff(device, static_cast<const rl_tools::nn::layers::embedding::LayerForward<SPEC_1>&>(l1), static_cast<const rl_tools::nn::layers::embedding::LayerForward<SPEC_2>&>(l2));
+    }
+    template <typename DEVICE, typename SPEC_1, typename SPEC_2>
+    RL_TOOLS_FUNCTION_PLACEMENT typename SPEC_1::T abs_diff(DEVICE& device, const rl_tools::nn::layers::embedding::LayerGradient<SPEC_1>& l1, const rl_tools::nn::layers::embedding::LayerGradient<SPEC_2>& l2) {
+        typename SPEC_1::T diff = abs_diff(device, l1.output, l2.output);
+        diff += abs_diff(device, static_cast<const rl_tools::nn::layers::embedding::LayerBackward<SPEC_1>&>(l1), static_cast<const rl_tools::nn::layers::embedding::LayerBackward<SPEC_2>&>(l2));
+        return diff;
+    }
+    template <typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void reset_forward_state(DEVICE& device, rl_tools::nn::layers::embedding::LayerForward<SPEC>& l) { }
+    template <typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void reset_forward_state(DEVICE& device, rl_tools::nn::layers::embedding::LayerBackward<SPEC>& l) {
+        reset_forward_state(device, (rl_tools::nn::layers::embedding::LayerForward<SPEC>&) l);
+    }
+    template <typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void reset_forward_state(DEVICE& device, rl_tools::nn::layers::embedding::LayerGradient<SPEC>& l) {
+        reset_forward_state(device, static_cast<rl_tools::nn::layers::embedding::LayerBackward<SPEC>&>(l));
+        set_all(device, l.output, (typename decltype(l.output)::SPEC::T)0);
+    }
+    template <typename DEVICE, typename SPEC, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT bool is_nan(DEVICE& device, const rl_tools::nn::layers::embedding::LayerForward<SPEC>& l, const Mode<MODE>& mode = Mode<mode::Default<>>{}) {
+        return is_nan(device, l.weights, mode);
+    }
+    template <typename DEVICE, typename SPEC, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT bool is_nan(DEVICE& device, const rl_tools::nn::layers::embedding::LayerBackward<SPEC>& l, const Mode<MODE>& mode = Mode<mode::Default<>>{}) {
+        return is_nan(device, static_cast<const rl_tools::nn::layers::embedding::LayerForward<SPEC>&>(l), mode);
+    }
+    template <typename DEVICE, typename SPEC, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT bool is_nan(DEVICE& device, const rl_tools::nn::layers::embedding::LayerGradient<SPEC>& l, const Mode<MODE>& mode = Mode<mode::Default<>>{}) {
+        bool upstream_nan = is_nan(device, static_cast<const rl_tools::nn::layers::embedding::LayerBackward<SPEC>&>(l), mode);
+        if constexpr(mode::is<MODE, nn::parameters::mode::ParametersOnly>){
+            return upstream_nan;
+        }
+        return upstream_nan || is_nan(device, l.output, mode);
+    }
+    template<typename DEVICE, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT bool is_nan(DEVICE& device, nn::layers::embedding::State& state, const Mode<MODE>& mode = Mode<mode::Default<>>{}){
+        return false;
+    }
+    template<typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT auto output(DEVICE& device, nn::layers::embedding::LayerGradient<SPEC>& l){
+        // return l.output;
+        auto tensor_flat = to_tensor(device, l.output);
+        auto tensor = view_memory<typename SPEC::OUTPUT_SHAPE>(device, tensor_flat);
+        return tensor;
+    }
+    template<typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT auto gradient_norm(DEVICE& device, const nn::layers::embedding::LayerGradient<SPEC>& layer) {
+        return gradient_norm(device, layer.weights);
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+#endif

hyperdrone/_vendor/rl-tools/include/rl_tools/nn/layers/embedding/persist.h

@@ -0,0 +1,46 @@
+#include "../../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_NN_LAYERS_EMBEDDING_PERSIST_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_NN_LAYERS_EMBEDDING_PERSIST_H
+
+#include "layer.h"
+#include <iostream>
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools {
+    template<typename DEVICE, typename SPEC, typename GROUP>
+    void save(DEVICE& device, nn::layers::embedding::LayerForward<SPEC>& layer, GROUP& group) {
+        // todo: forward implementation to Parameter struct
+        auto weights_group = create_group(device, group, "weights");
+        save(device, layer.weights, weights_group);
+        set_attribute(device, group, "type", "embedding");
+        write_attributes(device, group);
+    }
+    template<typename DEVICE, typename SPEC, typename GROUP>
+    void save(DEVICE& device, nn::layers::embedding::LayerBackward<SPEC>& layer, GROUP& group) {
+        save(device, (nn::layers::embedding::LayerForward<SPEC>&)layer, group);
+    }
+    template<typename DEVICE, typename SPEC, typename GROUP>
+    void save(DEVICE& device, nn::layers::embedding::LayerGradient<SPEC>& layer, GROUP& group) {
+        save(device, (nn::layers::embedding::LayerBackward<SPEC>&)layer, group);
+        save(device, layer.output, group, "output");
+    }
+    template<typename DEVICE, typename SPEC, typename GROUP>
+    bool load(DEVICE& device, nn::layers::embedding::LayerForward<SPEC>& layer, GROUP& group) {
+        auto weights_group = get_group(device, group, "weights");
+        return load(device, layer.weights, weights_group);
+    }
+    template<typename DEVICE, typename SPEC, typename GROUP>
+    bool load(DEVICE& device, nn::layers::embedding::LayerBackward<SPEC>& layer, GROUP& group) {
+        return load(device, (nn::layers::embedding::LayerForward<SPEC>&)layer, group);
+    }
+    template<typename DEVICE, typename SPEC, typename GROUP>
+    bool load(DEVICE& device, nn::layers::embedding::LayerGradient<SPEC>& layer, GROUP& group) {
+        bool success = load(device, (nn::layers::embedding::LayerBackward<SPEC>&)layer, group);
+        if(group_exists(device, group, "output")){
+            success &= load(device, layer.output, group, "output");
+        }
+        return success;
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+#endif

hyperdrone/_vendor/rl-tools/include/rl_tools/nn/layers/flatten/layer.h

@@ -0,0 +1,121 @@
+#include "../../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_NN_LAYERS_FLATTEN_LAYER_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_NN_LAYERS_FLATTEN_LAYER_H
+#include "../../../utils/generic/typing.h"
+#include "../../../containers/tensor/tensor.h"
+#include "../../../nn/capability/capability.h"
+#include "../../../nn/parameters/parameters.h"
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools::nn::layers::flatten {
+
+    template<typename T_TYPE_POLICY, typename T_TI>
+    struct Configuration{
+        using TYPE_POLICY = T_TYPE_POLICY;
+        using TI = T_TI;
+    };
+
+    template <typename T_CONFIG, typename T_CAPABILITY, typename T_INPUT_SHAPE>
+    struct Specification: T_CAPABILITY, T_CONFIG{
+        using CONFIG = T_CONFIG;
+        using TYPE_POLICY = typename CONFIG::TYPE_POLICY;
+        using TI = typename CONFIG::TI;
+        using CAPABILITY = T_CAPABILITY;
+        using INPUT_SHAPE = T_INPUT_SHAPE;
+        static_assert(length(INPUT_SHAPE{}) >= 4, "Flatten input shape must have at least 4 dimensions (...BATCH x H x W x C)");
+        static constexpr TI INPUT_HEIGHT = get<length(INPUT_SHAPE{})-3>(INPUT_SHAPE{});
+        static constexpr TI INPUT_WIDTH = get<length(INPUT_SHAPE{})-2>(INPUT_SHAPE{});
+        static constexpr TI INPUT_CHANNELS = get_last(INPUT_SHAPE{});
+        static constexpr TI OUTPUT_DIM = INPUT_HEIGHT * INPUT_WIDTH * INPUT_CHANNELS;
+
+        using BATCH_SHAPE = tensor::PopBack<tensor::PopBack<tensor::PopBack<INPUT_SHAPE>>>;
+        static constexpr TI INTERNAL_BATCH_SIZE = get<0>(tensor::CumulativeProduct<BATCH_SHAPE>{});
+
+        // Output shape: (...BATCH, H*W*C) — flatten spatial dims and channels into one
+        template <typename NEW_INPUT_SHAPE>
+        struct OUTPUT_SHAPE_FACTORY{
+            static_assert(length(NEW_INPUT_SHAPE{}) >= 4);
+            static constexpr TI NEW_H = get<length(NEW_INPUT_SHAPE{})-3>(NEW_INPUT_SHAPE{});
+            static constexpr TI NEW_W = get<length(NEW_INPUT_SHAPE{})-2>(NEW_INPUT_SHAPE{});
+            static constexpr TI NEW_C = get_last(NEW_INPUT_SHAPE{});
+            static_assert(NEW_H == INPUT_HEIGHT && NEW_W == INPUT_WIDTH && NEW_C == INPUT_CHANNELS);
+            using NEW_BATCH_SHAPE = tensor::PopBack<tensor::PopBack<tensor::PopBack<NEW_INPUT_SHAPE>>>;
+            using SHAPE = tensor::Append<NEW_BATCH_SHAPE, NEW_H * NEW_W * NEW_C>;
+        };
+        using OUTPUT_SHAPE = typename OUTPUT_SHAPE_FACTORY<INPUT_SHAPE>::SHAPE;
+        static constexpr TI NUM_WEIGHTS = 0;
+    };
+
+    template<typename SPEC_1, typename SPEC_2>
+    constexpr bool check_spec_memory =
+        SPEC_1::INPUT_HEIGHT == SPEC_2::INPUT_HEIGHT
+        && SPEC_1::INPUT_WIDTH == SPEC_2::INPUT_WIDTH
+        && SPEC_1::INPUT_CHANNELS == SPEC_2::INPUT_CHANNELS;
+
+    template<typename SPEC_1, typename SPEC_2>
+    constexpr bool check_spec = check_spec_memory<SPEC_1, SPEC_2>;
+
+    template <typename LAYER_SPEC, typename INPUT_SPEC, typename OUTPUT_SPEC>
+    constexpr bool check_input_output =
+        length(typename INPUT_SPEC::SHAPE{}) >= 4 &&
+        length(typename OUTPUT_SPEC::SHAPE{}) >= 2 &&
+        get<length(typename INPUT_SPEC::SHAPE{})-1>(typename INPUT_SPEC::SHAPE{}) == LAYER_SPEC::INPUT_CHANNELS &&
+        get<length(typename INPUT_SPEC::SHAPE{})-2>(typename INPUT_SPEC::SHAPE{}) == LAYER_SPEC::INPUT_WIDTH &&
+        get<length(typename INPUT_SPEC::SHAPE{})-3>(typename INPUT_SPEC::SHAPE{}) == LAYER_SPEC::INPUT_HEIGHT &&
+        get_last(typename OUTPUT_SPEC::SHAPE{}) == LAYER_SPEC::OUTPUT_DIM;
+
+    struct State{};
+    struct Buffer{};
+
+    template<typename T_SPEC>
+    struct LayerForward {
+        using SPEC = T_SPEC;
+        using TYPE_POLICY = typename SPEC::TYPE_POLICY;
+        using TI = typename SPEC::TI;
+        static constexpr TI INPUT_HEIGHT = SPEC::INPUT_HEIGHT;
+        static constexpr TI INPUT_WIDTH = SPEC::INPUT_WIDTH;
+        static constexpr TI INPUT_CHANNELS = SPEC::INPUT_CHANNELS;
+        static constexpr TI OUTPUT_DIM = SPEC::OUTPUT_DIM;
+        static constexpr TI NUM_WEIGHTS = SPEC::NUM_WEIGHTS;
+        static constexpr TI INTERNAL_BATCH_SIZE = SPEC::INTERNAL_BATCH_SIZE;
+        using INPUT_SHAPE = typename SPEC::INPUT_SHAPE;
+        template <typename NEW_INPUT_SHAPE>
+        using OUTPUT_SHAPE_FACTORY = typename SPEC::template OUTPUT_SHAPE_FACTORY<NEW_INPUT_SHAPE>::SHAPE;
+        using OUTPUT_SHAPE = typename SPEC::OUTPUT_SHAPE;
+        template<bool DYNAMIC_ALLOCATION=true>
+        using Buffer = flatten::Buffer;
+        template<bool DYNAMIC_ALLOCATION=true>
+        using State = flatten::State;
+    };
+
+    template<typename SPEC>
+    struct LayerBackward: public LayerForward<SPEC>{};
+
+    template<typename SPEC>
+    struct LayerGradient: public LayerBackward<SPEC>{
+        using T = typename SPEC::TYPE_POLICY::template GET<numeric_types::categories::Activation>;
+        using TI = typename SPEC::TI;
+        using OUTPUT_CONTAINER_SHAPE = tensor::Shape<TI, SPEC::INTERNAL_BATCH_SIZE, SPEC::OUTPUT_DIM>;
+        using OUTPUT_CONTAINER_SPEC = tensor::Specification<T, TI, OUTPUT_CONTAINER_SHAPE, SPEC::DYNAMIC_ALLOCATION, tensor::RowMajorStride<OUTPUT_CONTAINER_SHAPE>, SPEC::CONST>;
+        using OUTPUT_CONTAINER_TYPE = Tensor<OUTPUT_CONTAINER_SPEC>;
+        OUTPUT_CONTAINER_TYPE output;
+    };
+
+    template<typename CONFIG, typename CAPABILITY, typename INPUT_SHAPE>
+    using Layer =
+        typename utils::typing::conditional_t<CAPABILITY::TAG == nn::LayerCapability::Forward,
+            LayerForward<Specification<CONFIG, CAPABILITY, INPUT_SHAPE>>,
+        typename utils::typing::conditional_t<CAPABILITY::TAG == nn::LayerCapability::Backward,
+            LayerBackward<Specification<CONFIG, CAPABILITY, INPUT_SHAPE>>,
+        typename utils::typing::conditional_t<CAPABILITY::TAG == nn::LayerCapability::Gradient,
+            LayerGradient<Specification<CONFIG, CAPABILITY, INPUT_SHAPE>>, void>>>;
+
+    template <typename CONFIG>
+    struct BindConfiguration{
+        template <typename CAPABILITY, typename INPUT_SHAPE>
+        using Layer = nn::layers::flatten::Layer<CONFIG, CAPABILITY, INPUT_SHAPE>;
+    };
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+#endif