hyperray 0.0.1__py3-none-any.whl

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

@@ -0,0 +1,295 @@
+#include "../../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_NN_LAYERS_STANDARDIZE_OPERATIONS_GENERIC_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_NN_LAYERS_STANDARDIZE_OPERATIONS_GENERIC_H
+#include "../../../utils/generic/typing.h"
+#include "../../../containers/matrix/matrix.h"
+
+#include "../../../nn/nn.h"
+#include "../../../nn/parameters/parameters.h"
+#include "layer.h"
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools{
+    template <typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void malloc(DEVICE& device, nn::layers::standardize::LayerForward<SPEC>& layer){
+        malloc(device, layer.mean);
+        malloc(device, layer.precision);
+        malloc(device, layer.running_mean);
+        malloc(device, layer.running_std);
+    }
+    template <typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void free(DEVICE& device, nn::layers::standardize::LayerForward<SPEC>& layer){
+        free(device, layer.mean);
+        free(device, layer.precision);
+        free(device, layer.running_mean);
+        free(device, layer.running_std);
+    }
+    template <typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void malloc(DEVICE& device, nn::layers::standardize::LayerGradient<SPEC>& layer){
+        malloc(device, static_cast<nn::layers::standardize::LayerForward<SPEC>&>(layer));
+        malloc(device, layer.output);
+    }
+    template <typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void free(DEVICE& device, nn::layers::standardize::LayerGradient<SPEC>& layer){
+        free(device, static_cast<nn::layers::standardize::LayerForward<SPEC>&>(layer));
+        free(device, layer.output);
+    }
+    template<typename DEVICE>
+    RL_TOOLS_FUNCTION_PLACEMENT void malloc(DEVICE& device, nn::layers::standardize::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::standardize::State& source, nn::layers::standardize::State& target){}
+    template<typename DEVICE, typename SPEC, typename RNG, typename MODE>
+    RL_TOOLS_FUNCTION_PLACEMENT void reset(DEVICE& device, const nn::layers::standardize::LayerForward<SPEC>& layer, nn::layers::standardize::State& state, RNG&, Mode<MODE> mode = Mode<mode::Default<>>{}) { } // no-op
+    template<typename DEVICE>
+    RL_TOOLS_FUNCTION_PLACEMENT void free(DEVICE& device, nn::layers::standardize::State& state) { } // no-op
+    template <typename DEVICE>
+    RL_TOOLS_FUNCTION_PLACEMENT void malloc(DEVICE& device, nn::layers::standardize::Buffer& buffer){ }
+    template <typename DEVICE>
+    RL_TOOLS_FUNCTION_PLACEMENT void free(DEVICE& device, nn::layers::standardize::Buffer& buffer){ }
+    template<typename DEVICE, typename SPEC, typename RNG>
+    RL_TOOLS_FUNCTION_PLACEMENT void init_weights(DEVICE& device, nn::layers::standardize::LayerForward<SPEC>& layer, RNG& rng) {
+        set_all(device, layer.mean.parameters, 0);
+        set_all(device, layer.precision.parameters, 1);
+        layer.age = 0;
+        set_all(device, layer.running_mean.parameters, 0);
+        set_all(device, layer.running_std.parameters, 1);
+    }
+    namespace nn::layers::standardize{
+        template<typename DEVICE, typename LAYER_SPEC, typename INPUT_SPEC>
+        RL_TOOLS_FUNCTION_PLACEMENT void _accumulate(DEVICE& device, nn::layers::standardize::LayerForward<LAYER_SPEC>& layer, const Matrix<INPUT_SPEC>& data) {
+            static_assert(INPUT_SPEC::COLS == LAYER_SPEC::INPUT_DIM, "Data dimension must match layer input dimension");
+            using T = typename LAYER_SPEC::TYPE_POLICY::DEFAULT;
+            using TI = typename DEVICE::index_t;
+            constexpr TI DIM = LAYER_SPEC::INPUT_DIM;
+            constexpr TI BATCH_SIZE = INPUT_SPEC::ROWS;
+            static_assert(BATCH_SIZE > 1, "Batch size must be greater than 1");
+            layer.age++;
+            for(TI dim_i = 0; dim_i < DIM; dim_i++){
+                auto column = col(device, data, dim_i);
+                T data_mean = mean(device, column);
+                T data_std = std(device, column);
+                T current_mean = get(device, layer.running_mean.parameters, dim_i);
+                T new_mean = current_mean + (data_mean - current_mean) / layer.age;
+                T current_std = get(device, layer.running_std.parameters, dim_i);
+                T new_std = current_std + (data_std - current_std) / layer.age;
+                set(device, layer.running_mean.parameters, new_mean, dim_i);
+                set(device, layer.running_std.parameters, new_std, dim_i);
+                set(device, layer.mean.parameters, new_mean, dim_i);
+                T precision = new_std == 0 ? 0 : 1 / new_std;
+                set(device, layer.precision.parameters, precision, dim_i);
+            }
+        }
+    }
+    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::standardize::LayerForward<LAYER_SPEC>& layer, const Matrix<INPUT_SPEC>& input, Matrix<OUTPUT_SPEC>& output, nn::layers::standardize::Buffer& buffer, RNG& rng, const Mode<MODE>& mode = Mode<mode::Default<>>{}) {
+        static_assert(nn::layers::standardize::check_input_output<LAYER_SPEC, INPUT_SPEC, OUTPUT_SPEC>);
+        using T = typename LAYER_SPEC::TYPE_POLICY::DEFAULT;
+        using TI = typename DEVICE::index_t;
+        constexpr TI BATCH_SIZE = INPUT_SPEC::ROWS;
+        for(TI batch_i=0; batch_i < BATCH_SIZE; batch_i++){
+            for(TI output_i = 0; output_i < LAYER_SPEC::OUTPUT_DIM; output_i++) {
+                T mean = get(device, layer.mean.parameters, output_i);
+                T precision = get(device, layer.precision.parameters, output_i);
+                T input_value = get(input, batch_i, output_i);
+                T output_value = (input_value - mean);
+                if(precision != 0){
+                    output_value *= precision;
+                }
+                set(output, batch_i, output_i, output_value);
+            }
+        }
+    }
+    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::standardize::LayerBackward<LAYER_SPEC>& layer, const Matrix<INPUT_SPEC>& input, Matrix<OUTPUT_SPEC>& output, nn::layers::standardize::Buffer& buffer, RNG& rng, const Mode<MODE>& mode = Mode<mode::Default<>>{}){
+        static_assert(nn::layers::standardize::check_input_output<LAYER_SPEC, INPUT_SPEC, OUTPUT_SPEC>);
+        if constexpr(mode::is<MODE, nn::layers::standardize::AccumulateMode>){
+            nn::layers::standardize::_accumulate(device, static_cast<nn::layers::standardize::LayerForward<LAYER_SPEC>&>(layer), input);
+        }
+        evaluate(device, layer, input, output, buffer, rng, mode);
+    }
+    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::standardize::LayerGradient<LAYER_SPEC>& layer, const Matrix<INPUT_SPEC>& input, nn::layers::standardize::Buffer& buffer, RNG& rng, const Mode<MODE>& mode = Mode<mode::Default<>>{}){
+        static_assert(nn::layers::standardize::check_input_output<LAYER_SPEC, INPUT_SPEC, typename decltype(layer.output)::SPEC>);
+        forward(device, layer, input, layer.output, buffer, rng, mode);
+    }
+
+    template<typename DEVICE, typename LAYER_SPEC, typename D_OUTPUT_SPEC, typename D_INPUT_SPEC, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT void backward_input(DEVICE& device, nn::layers::standardize::LayerBackward<LAYER_SPEC>& layer, const Matrix<D_OUTPUT_SPEC>& d_output, Matrix<D_INPUT_SPEC>& d_input, nn::layers::standardize::Buffer& buffer, const Mode<MODE>& mode = Mode<mode::Default<>>{}){
+        static_assert(nn::layers::standardize::check_input_output<LAYER_SPEC, D_INPUT_SPEC, D_OUTPUT_SPEC>);
+        static_assert(LAYER_SPEC::INPUT_DIM == LAYER_SPEC::OUTPUT_DIM);
+        using TI = typename DEVICE::index_t;
+        constexpr TI OUTPUT_DIM = LAYER_SPEC::OUTPUT_DIM;
+        constexpr TI BATCH_SIZE = D_OUTPUT_SPEC::ROWS;
+        using T = typename LAYER_SPEC::TYPE_POLICY::DEFAULT;
+
+        for(TI batch_i=0; batch_i < BATCH_SIZE; batch_i++){
+            for(TI output_i = 0; output_i < OUTPUT_DIM; output_i++) {
+                T d_output_value = get(d_output, batch_i, output_i);
+                T precision = get(device, layer.precision.parameters, output_i);
+                T d_input_value = d_output_value * precision;
+                set(d_input, batch_i, output_i, d_input_value);
+            }
+        }
+    }
+
+    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::standardize::LayerGradient<LAYER_SPEC>& layer, const Matrix<INPUT_SPEC>& input, Matrix<D_OUTPUT_SPEC>& d_output, nn::layers::standardize::Buffer& buffer, const Mode<MODE>& mode = Mode<mode::Default<>>{}) {
+        // this is a no-op as the standardize layer does not have trainable parameters
+    }
+
+    template<typename DEVICE, typename LAYER_SPEC, typename INPUT_SPEC, typename D_OUTPUT_SPEC, typename D_INPUT_SPEC, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT void backward_full(DEVICE& device, nn::layers::standardize::LayerGradient<LAYER_SPEC>& layer, const Matrix<INPUT_SPEC>& input, Matrix<D_OUTPUT_SPEC>& d_output, Matrix<D_INPUT_SPEC>& d_input, nn::layers::standardize::Buffer& buffer, const Mode<MODE>& mode = Mode<mode::Default<>>{}) {
+        backward_input(device, layer, d_output, d_input, buffer, mode);
+    }
+    template<typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void zero_gradient(DEVICE& device, nn::layers::standardize::LayerGradient<SPEC>& layer) {
+    }
+    template<typename DEVICE, typename SOURCE_SPEC, typename TARGET_SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void add_gradient(DEVICE& device, nn::layers::standardize::LayerGradient<SOURCE_SPEC>& source, nn::layers::standardize::LayerGradient<TARGET_SPEC>& target) {}
+    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::standardize::LayerGradient<SOURCE_SPEC>& source, nn::layers::standardize::LayerGradient<TARGET_SPEC>& target) {}
+    template<typename DEVICE, typename SPEC, typename OPTIMIZER>
+    RL_TOOLS_FUNCTION_PLACEMENT void update(DEVICE& device, nn::layers::standardize::LayerGradient<SPEC>& layer, OPTIMIZER& optimizer){
+    }
+
+    template <typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void reset_forward_state(DEVICE& device, nn::layers::standardize::LayerBackward<SPEC>& l) { }
+    template <typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void reset_forward_state(DEVICE& device, nn::layers::standardize::LayerGradient<SPEC>& l) {
+        reset_forward_state(device, static_cast<nn::layers::standardize::LayerBackward<SPEC>&>(l));
+        set_all(device, l.output, 0);
+    }
+
+    template<typename DEVICE, typename SPEC, typename OPTIMIZER>
+    RL_TOOLS_FUNCTION_PLACEMENT void _reset_optimizer_state(DEVICE& device, nn::layers::standardize::LayerGradient<SPEC>& layer, OPTIMIZER& 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::standardize::LayerForward<SOURCE_SPEC>& source, nn::layers::standardize::LayerForward<TARGET_SPEC>& target){
+        static_assert(nn::layers::standardize::check_compatibility<SOURCE_SPEC, TARGET_SPEC>);
+        copy(source_device, target_device, source.mean, target.mean);
+        copy(source_device, target_device, source.precision, target.precision);
+        target.age = source.age;
+        copy(source_device, target_device, source.running_mean, target.running_mean);
+        copy(source_device, target_device, source.running_std, target.running_std);
+    }
+    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::standardize::LayerGradient<SOURCE_SPEC>& source, nn::layers::standardize::LayerGradient<TARGET_SPEC>& target){
+        static_assert(nn::layers::standardize::check_compatibility<SOURCE_SPEC, TARGET_SPEC>);
+        copy(source_device, target_device, static_cast<const nn::layers::standardize::LayerForward<SOURCE_SPEC>&>(source), static_cast<nn::layers::standardize::LayerForward<TARGET_SPEC>&>(target));
+        copy(source_device, target_device, source.output, target.output);
+    }
+    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::standardize::LayerForward<SPEC_1>& l1, const rl_tools::nn::layers::standardize::LayerForward<SPEC_2>& l2) {
+        static_assert(nn::layers::standardize::check_compatibility<SPEC_1, SPEC_2>);
+        using T = typename SPEC_1::TYPE_POLICY::DEFAULT;
+        T acc = 0;
+        acc += abs_diff(device, l1.mean, l2.mean);
+        acc += abs_diff(device, l1.precision, l2.precision);
+        acc += math::abs(device.math, (T)l1.age - (T)l2.age);
+        acc += abs_diff(device, l1.running_mean, l2.running_mean);
+        acc += abs_diff(device, l1.running_std, l2.running_std);
+        return acc;
+    }
+    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::standardize::LayerGradient<SPEC_1>& l1, const rl_tools::nn::layers::standardize::LayerGradient<SPEC_2>& l2) {
+        static_assert(nn::layers::standardize::check_compatibility<SPEC_1, SPEC_2>);
+        return abs_diff(device, static_cast<const rl_tools::nn::layers::standardize::LayerForward<SPEC_1>&>(l1), static_cast<const rl_tools::nn::layers::standardize::LayerForward<SPEC_2>&>(l2));
+    }
+    template <typename DEVICE>
+    RL_TOOLS_FUNCTION_PLACEMENT auto abs_diff(DEVICE& device, const rl_tools::nn::layers::standardize::State& s1, const rl_tools::nn::layers::standardize::State& s2) {
+        return 0;
+    }
+    template<typename DEVICE, typename LAYER_SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT constexpr auto output(DEVICE& device, nn::layers::standardize::LayerGradient<LAYER_SPEC>& l){
+        auto tensor_flat = to_tensor(device, l.output);
+        auto tensor = view_memory<typename LAYER_SPEC::OUTPUT_SHAPE>(device, tensor_flat);
+        return tensor;
+    }
+    template <typename DEVICE, typename SPEC, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT bool is_nan(DEVICE& device, const rl_tools::nn::layers::standardize::LayerForward<SPEC>& l, const Mode<MODE>& mode = Mode<mode::Default<>>{}) {
+        if constexpr(mode::is<MODE, nn::parameters::mode::ParametersOnly>){
+            return false;
+        }
+        return is_nan(device, l.mean, mode) || is_nan(device, l.precision, mode) || is_nan(device, l.running_mean, mode) || is_nan(device, l.running_std, mode);
+    }
+    template <typename DEVICE, typename SPEC, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT bool is_nan(DEVICE& device, const rl_tools::nn::layers::standardize::LayerBackward<SPEC>& l, const Mode<MODE>& mode = Mode<mode::Default<>>{}) {
+        return is_nan(device, static_cast<const rl_tools::nn::layers::standardize::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::standardize::LayerGradient<SPEC>& l, const Mode<MODE>& mode = Mode<mode::Default<>>{}) {
+        bool upstream_nan = is_nan(device, static_cast<const rl_tools::nn::layers::standardize::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::standardize::State& state, const Mode<MODE>& mode = Mode<mode::Default<>>{}){
+        return false;
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+
+// Tensor proxies
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools{
+    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::standardize::LayerForward<LAYER_SPEC>& layer, const Tensor<INPUT_SPEC>& input, Tensor<OUTPUT_SPEC>& output, nn::layers::standardize::Buffer& buffer, RNG& rng, const Mode<MODE>& mode = Mode<mode::Default<>>{}) {
+        auto matrix_view_input = matrix_view(device, input);
+        auto matrix_view_output = matrix_view(device, output);
+        evaluate(device, layer, matrix_view_input, matrix_view_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 evaluate_step(DEVICE& device, const nn::layers::standardize::LayerForward<LAYER_SPEC>& layer, const Tensor<INPUT_SPEC>& input, nn::layers::standardize::State& state, Tensor<OUTPUT_SPEC>& output, nn::layers::standardize::Buffer& buffer, RNG& rng, const Mode<MODE>& mode = Mode<mode::Default<>>{}) {
+        auto matrix_view_input = matrix_view(device, input);
+        auto matrix_view_output = matrix_view(device, output);
+        evaluate(device, layer, matrix_view_input, matrix_view_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::standardize::LayerBackward<LAYER_SPEC>& layer, const Tensor<INPUT_SPEC>& input, Tensor<OUTPUT_SPEC>& output, nn::layers::standardize::Buffer& buffer, RNG& rng, const Mode<MODE>& mode = Mode<mode::Default<>>{}){
+        auto matrix_view_input = matrix_view(device, input);
+        auto matrix_view_output = matrix_view(device, output);
+        forward(device, layer, matrix_view_input, matrix_view_output, buffer, rng, mode);
+    }
+    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::standardize::LayerGradient<LAYER_SPEC>& layer, const Tensor<INPUT_SPEC>& input, nn::layers::standardize::Buffer& buffer, RNG& rng, const Mode<MODE>& mode = Mode<mode::Default<>>{}) {
+        auto matrix_view_input = matrix_view(device, input);
+        forward(device, layer, matrix_view_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::standardize::LayerGradient<LAYER_SPEC>& layer, const Tensor<INPUT_SPEC>& input, Tensor<OUTPUT_SPEC>& output, nn::layers::standardize::Buffer& buffer, RNG& rng, const Mode<MODE>& mode = Mode<mode::Default<>>{}) {
+        auto matrix_view_input = matrix_view(device, input);
+        auto matrix_view_output = matrix_view(device, output);
+        forward(device, layer, matrix_view_input, matrix_view_output, buffer, rng, mode);
+    }
+    template<typename DEVICE, typename LAYER_SPEC, typename D_OUTPUT_SPEC, typename D_INPUT_SPEC, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT void backward_input(DEVICE& device, const nn::layers::standardize::LayerBackward<LAYER_SPEC>& layer, const Tensor<D_OUTPUT_SPEC>& d_output, Tensor<D_INPUT_SPEC>& d_input, nn::layers::standardize::Buffer& buffer, const Mode<MODE>& mode = Mode<mode::Default<>>{}){
+        auto matrix_view_d_output = matrix_view(device, d_output);
+        auto matrix_view_d_input = matrix_view(device, d_input);
+        backward_input(device, layer, matrix_view_d_output, matrix_view_d_input, buffer, mode);
+    }
+
+    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::standardize::LayerGradient<LAYER_SPEC>& layer, const Tensor<INPUT_SPEC>& input, Tensor<D_OUTPUT_SPEC>& d_output, nn::layers::standardize::Buffer& buffer, const Mode<MODE>& mode = Mode<mode::Default<>>{}) {
+        auto matrix_view_input = matrix_view(device, input);
+        auto matrix_view_d_output = matrix_view(device, d_output);
+        backward(device, layer, matrix_view_input, matrix_view_d_output, buffer, mode);
+    }
+
+    template<typename DEVICE, typename LAYER_SPEC, typename INPUT_SPEC, typename D_OUTPUT_SPEC, typename D_INPUT_SPEC, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT void backward_full(DEVICE& device, nn::layers::standardize::LayerGradient<LAYER_SPEC>& layer, const Tensor<INPUT_SPEC>& input, Tensor<D_OUTPUT_SPEC>& d_output, Tensor<D_INPUT_SPEC>& d_input, nn::layers::standardize::Buffer& buffer, const Mode<MODE>& mode = Mode<mode::Default<>>{}) {
+        auto matrix_view_input = matrix_view(device, input);
+        auto matrix_view_d_output = matrix_view(device, d_output);
+        auto matrix_view_d_input = matrix_view(device, d_input);
+        backward_full(device, layer, matrix_view_input, matrix_view_d_output, matrix_view_d_input, buffer, mode);
+    }
+    template<typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT auto gradient_norm(DEVICE& device, const nn::layers::standardize::LayerGradient<SPEC>& layer) {
+        return 0;
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+
+#endif

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

@@ -0,0 +1,83 @@
+#include "../../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_NN_LAYERS_STANDARDIZE_PERSIST_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_NN_LAYERS_STANDARDIZE_PERSIST_H
+#include "layer.h"
+#include "../../parameters/persist.h"
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools {
+    template<typename DEVICE, typename SPEC, typename GROUP>
+    void save(DEVICE& device, nn::layers::standardize::LayerForward<SPEC>& layer, GROUP& group) {
+        auto mean_group = create_group(device, group, "mean");
+        save(device, layer.mean, mean_group);
+        auto precision_group = create_group(device, group, "precision");
+        save(device, layer.precision, precision_group);
+        auto running_mean_group = create_group(device, group, "running_mean");
+        save(device, layer.running_mean, running_mean_group);
+        write_attributes(device, running_mean_group);
+        auto running_std_group = create_group(device, group, "running_std");
+        save(device, layer.running_std, running_std_group);
+        write_attributes(device, running_std_group);
+        {
+            using T = typename SPEC::TYPE_POLICY::DEFAULT;
+            auto age_group = create_group(device, group, "age");
+            Matrix<matrix::Specification<T, typename SPEC::TI, 1, 1, false>> age_matrix;
+            set(age_matrix, 0, 0, layer.age);
+            save(device, age_matrix, age_group, "value");
+            write_attributes(device, age_group);
+        }
+        set_attribute(device, group, "type", "standardize");
+        write_attributes(device, group);
+    }
+    template<typename DEVICE, typename SPEC, typename GROUP>
+    void save(DEVICE& device, nn::layers::standardize::LayerBackward<SPEC>& layer, GROUP& group) {
+        save(device, (nn::layers::standardize::LayerForward<SPEC>&)layer, group);
+    }
+    template<typename DEVICE, typename SPEC, typename GROUP>
+    void save(DEVICE& device, nn::layers::standardize::LayerGradient<SPEC>& layer, GROUP& group) {
+        save(device, (nn::layers::standardize::LayerBackward<SPEC>&)layer, group);
+        save(device, layer.output, group, "output");
+    }
+    template<typename DEVICE, typename SPEC, typename GROUP>
+    bool load(DEVICE& device, nn::layers::standardize::LayerForward<SPEC>& layer, GROUP& group) {
+        auto mean_group = get_group(device, group, "mean");
+        bool success = load(device, layer.mean, mean_group);
+        auto precision_group = get_group(device, group, "precision");
+        success &= load(device, layer.precision, precision_group);
+        if(group_exists(device, group, "running_mean")){
+            auto running_mean_group = get_group(device, group, "running_mean");
+            success &= load(device, layer.running_mean, running_mean_group);
+        }
+        if(group_exists(device, group, "running_std")){
+            auto running_std_group = get_group(device, group, "running_std");
+            success &= load(device, layer.running_std, running_std_group);
+        }
+        if(group_exists(device, group, "age")){
+            using T = typename SPEC::TYPE_POLICY::DEFAULT;
+            auto age_group = get_group(device, group, "age");
+            Matrix<matrix::Specification<T, typename SPEC::TI, 1, 1, false>> age_matrix;
+            set(age_matrix, 0, 0, (T)0);
+            load(device, age_matrix, age_group, "value");
+            layer.age = get(age_matrix, 0, 0);
+        }
+        return success;
+    }
+    template<typename DEVICE, typename SPEC, typename GROUP>
+    bool load(DEVICE& device, nn::layers::standardize::LayerBackward<SPEC>& layer, GROUP& group) {
+        return load(device, (nn::layers::standardize::LayerForward<SPEC>&)layer, group);
+    }
+    template<typename DEVICE, typename SPEC, typename GROUP>
+    bool load(DEVICE& device, nn::layers::standardize::LayerGradient<SPEC>& layer, GROUP& group) {
+        bool success = load(device, (nn::layers::standardize::LayerBackward<SPEC>&)layer, group);
+        if(group_exists(device, group, "output")){
+            success &=  load(device, layer.output, group, "output");
+        }
+        return success;
+    }
+    template<typename DEVICE, typename GROUP>
+    void save(DEVICE& device, nn::layers::standardize::State& state, GROUP& group) {}
+    template<typename DEVICE, typename GROUP>
+    bool load(DEVICE& device, nn::layers::standardize::State& state, GROUP& group) { return true; }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+#endif

hyperdrone/_vendor/rl-tools/include/rl_tools/nn/layers/standardize/persist_code.h

@@ -0,0 +1,163 @@
+#include "../../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_NN_LAYERS_STANDARDIZE_PERSIST_CODE_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_NN_LAYERS_STANDARDIZE_PERSIST_CODE_H
+#include "layer.h"
+#include "../../../containers/matrix/persist_code.h"
+#include <sstream>
+#include "../../../persist/code.h"
+#include "../../../containers/matrix/persist_code.h"
+#include "../../../nn/capability/persist_code.h"
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools {
+    namespace nn::layers::standardize::persist_code{
+        template<typename DEVICE, typename SPEC>
+        rl_tools::persist::Code finish(DEVICE& device, nn::layers::standardize::LayerForward<SPEC> &layer, std::string name, rl_tools::persist::Code input, bool const_declaration=true, typename DEVICE::index_t indent=0){
+            using TI = typename DEVICE::index_t;
+            std::stringstream indent_ss;
+            for(TI i=0; i < indent; i++){
+                indent_ss << "    ";
+            }
+            std::string ind = indent_ss.str();
+            using TI = typename DEVICE::index_t;
+            std::string T_string = containers::persist::get_type_string<typename SPEC::TYPE_POLICY::DEFAULT>();
+            std::string TI_string = containers::persist::get_type_string<typename SPEC::TI>();
+            std::stringstream ss, ss_header;
+            ss_header << input.header;
+            ss_header << "#include <rl_tools/nn/layers/standardize/layer.h>\n";
+            ss << input.body;
+            ss << ind << "namespace " << name << " {\n";
+            ss << ind << "    using TYPE_POLICY = " << to_string(typename SPEC::TYPE_POLICY{}) << ";\n";
+            ss << ind << "    using CONFIG = " << "RL_TOOLS""_NAMESPACE_WRAPPER ::rl_tools::nn::layers::standardize::Configuration<"
+               << "TYPE_POLICY, "
+               << TI_string
+               << ">; \n";
+            ss << ind << "    " << "using TEMPLATE = RL_TOOLS""_NAMESPACE_WRAPPER ::rl_tools::nn::layers::standardize::BindConfiguration<CONFIG>;" << "\n";
+            ss << ind << "    " << "using INPUT_SHAPE = RL_TOOLS""_NAMESPACE_WRAPPER ::rl_tools::tensor::Shape<" << TI_string << ", " << get<0>(typename SPEC::INPUT_SHAPE{}) << ", " << get<1>(typename SPEC::INPUT_SHAPE{}) << ", " << get<2>(typename SPEC::INPUT_SHAPE{}) << ">;\n";
+            ss << ind << "    " << "using CAPABILITY = " << to_string(typename SPEC::CAPABILITY::template CHANGE_PARAMETERS<true, true>{}) << ";" << "\n";
+            ss << ind << "    " << "using TYPE = RL_TOOLS""_NAMESPACE_WRAPPER ::rl_tools::nn::layers::standardize::Layer<CONFIG, CAPABILITY, INPUT_SHAPE>;" << "\n";
+            std::string initializer_list;
+            if constexpr(SPEC::CAPABILITY::TAG == nn::LayerCapability::Forward){
+                initializer_list = "{mean::parameters, precision::parameters, age, running_mean::parameters, running_std::parameters}";
+            }
+            else{
+                if constexpr(SPEC::CAPABILITY::TAG == nn::LayerCapability::Backward){
+                    initializer_list = "{{mean::parameters, precision::parameters, age, running_mean::parameters, running_std::parameters}}";
+                }
+                else{
+                    if constexpr(SPEC::CAPABILITY::TAG == nn::LayerCapability::Gradient){
+                        initializer_list = "{{{mean::parameters, precision::parameters, age, running_mean::parameters, running_std::parameters}}, output::container}";
+                    }
+                    else{
+                        utils::assert_exit(device, false, "Unknown capability");
+                    }
+                }
+            }
+            ss << ind << "    " << (const_declaration ? "constexpr " : "") << "TYPE module = " << initializer_list << ";\n";
+            ss << ind << "    " << "template <typename T_TYPE = TYPE>" << "\n";
+            ss << ind << "    " << (const_declaration ? "constexpr " : "") << "T_TYPE factory = " << initializer_list << ";" << "\n";
+            ss << ind << "    " << "template <typename T_TYPE = TYPE>" << "\n";
+            ss << ind << "    " << (const_declaration ? "constexpr " : "") << "T_TYPE factory_function(){return T_TYPE" << initializer_list << ";" << "}\n";
+            ss << ind << "}\n";
+
+
+            return {ss_header.str(), ss.str()};
+        }
+    }
+    template<typename DEVICE, typename SPEC>
+    persist::Code save_code_split(DEVICE& device, nn::layers::standardize::LayerForward<SPEC> &layer, std::string name, bool const_declaration=true, typename DEVICE::index_t indent=0, bool finish=true){
+        using TI = typename DEVICE::index_t;
+        std::stringstream indent_ss;
+        for(TI i=0; i < indent; i++){
+            indent_ss << "    ";
+        }
+        std::string ind = indent_ss.str();
+        using TI = typename DEVICE::index_t;
+        std::stringstream ss, ss_header;
+        ss << ind << "namespace " << name << " {\n";
+        auto mean = save_code_split(device, layer.mean, "mean", const_declaration, indent+1);
+        ss_header << mean.header;
+        ss << mean.body;
+        auto precision = save_code_split(device, layer.precision, "precision", const_declaration, indent+1);
+        ss_header << precision.header;
+        ss << precision.body;
+        {
+            std::string T_string = containers::persist::get_type_string<typename SPEC::TYPE_POLICY::DEFAULT>();
+            ss << ind << "    " << (const_declaration ? "constexpr " : "") << T_string << " age = " << layer.age << ";\n";
+        }
+        auto running_mean = save_code_split(device, layer.running_mean, "running_mean", const_declaration, indent+1);
+        ss_header << running_mean.header;
+        ss << running_mean.body;
+        auto running_std = save_code_split(device, layer.running_std, "running_std", const_declaration, indent+1);
+        ss_header << running_std.header;
+        ss << running_std.body;
+        ss << ind << "}\n";
+        if(finish){
+            return nn::layers::standardize::persist_code::finish(device, layer, name, {ss_header.str(), ss.str()}, const_declaration, indent);
+        }
+        else{
+            return {ss_header.str(), ss.str()};
+        }
+    }
+    template<typename DEVICE, typename SPEC>
+    persist::Code save_code_split(DEVICE& device, nn::layers::standardize::LayerBackward<SPEC> &layer, std::string name, bool const_declaration=true, typename DEVICE::index_t indent=0, bool finish=true){
+        using TI = typename DEVICE::index_t;
+        std::stringstream indent_ss;
+        for(TI i=0; i < indent; i++){
+            indent_ss << "    ";
+        }
+        std::string ind = indent_ss.str();
+        using TI = typename DEVICE::index_t;
+        std::stringstream ss, ss_header;
+        auto previous = save_code_split(device, static_cast<nn::layers::standardize::LayerForward<SPEC>&>(layer), name, const_declaration, indent, false);
+        ss_header << previous.header;
+        ss << previous.body;
+        if(finish){
+            return nn::layers::standardize::persist_code::finish(device, layer, name, {ss_header.str(), ss.str()}, const_declaration, indent);
+        }
+        else{
+            return {ss_header.str(), ss.str()};
+        }
+    }
+
+    template<typename DEVICE, typename SPEC>
+    persist::Code save_code_split(DEVICE& device, nn::layers::standardize::LayerGradient<SPEC> &layer, std::string name, bool const_declaration=true, typename DEVICE::index_t indent=0){
+        using TI = typename DEVICE::index_t;
+        std::stringstream indent_ss;
+        for(TI i=0; i < indent; i++){
+            indent_ss << "    ";
+        }
+        std::string ind = indent_ss.str();
+        using TI = typename DEVICE::index_t;
+        std::stringstream ss, ss_header;
+        auto previous = save_code_split(device, static_cast<nn::layers::standardize::LayerBackward<SPEC>&>(layer), name, const_declaration, indent, false);
+        ss_header << previous.header;
+        ss << previous.body;
+        ss << ind << "namespace " << name << " {\n";
+        auto output = save_code_split(device, layer.output, "output", const_declaration, indent+1);
+        ss_header << output.header;
+        ss << output.body;
+        ss << ind << "}\n";
+        return nn::layers::standardize::persist_code::finish(device, layer, name, {ss_header.str(), ss.str()}, const_declaration, indent);
+    }
+
+    template<typename DEVICE, typename SPEC>
+    std::string save_code(DEVICE& device, nn::layers::standardize::LayerForward<SPEC> &layer, std::string name, bool const_declaration=true, typename DEVICE::index_t indent=0){
+        auto code = save_code_split(device, layer, name, const_declaration, indent);
+        return code.header + code.body;
+    }
+    template <typename DEVICE, typename SPEC>
+    std::string nn_analytics(DEVICE& device, nn::layers::standardize::LayerGradient<SPEC>& layer) {
+        std::string data;
+        data += "{";
+        data += "\"mean\": " + nn_analytics(device, layer.mean) + ", ";
+        data += "\"precision\": " + nn_analytics(device, layer.precision) + ", ";
+        data += "\"output\": " + json(device, layer.output);
+        data += "}";
+        return data;
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+
+#endif