hyperray 0.0.1__py3-none-any.whl

hyperdrone/_vendor/rl-tools/include/rl_tools/nn_models/mlp/network.h

@@ -0,0 +1,174 @@
+#include "../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_NN_MODELS_MLP_NETWORK_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_NN_MODELS_MLP_NETWORK_H
+
+#include "../../nn/nn.h"
+#include "../../numeric_types/categories.h"
+#include "../../nn/parameters/parameters.h"
+#include "../../nn/layers/dense/layer.h"
+#include "../../nn/optimizers/sgd/sgd.h"
+#include "../../nn/optimizers/adam/adam.h"
+#include "../../utils/generic/typing.h"
+#include "../../containers/matrix/matrix.h"
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools::nn_models::mlp {
+    template <typename T_TYPE_POLICY, typename T_TI, T_TI T_OUTPUT_DIM, T_TI T_NUM_LAYERS, T_TI T_HIDDEN_DIM, nn::activation_functions::ActivationFunction T_HIDDEN_ACTIVATION_FUNCTION, nn::activation_functions::ActivationFunction T_OUTPUT_ACTIVATION_FUNCTION, typename T_LAYER_INITIALIZER=nn::layers::dense::DefaultInitializer<T_TYPE_POLICY, T_TI>, bool T_IS_INPUT=true, bool T_IS_OUTPUT=true>
+    struct Configuration{
+        using TYPE_POLICY = T_TYPE_POLICY;
+        using TI = T_TI;
+        static constexpr T_TI OUTPUT_DIM = T_OUTPUT_DIM;
+        static constexpr T_TI NUM_LAYERS = T_NUM_LAYERS; // The input and output layers count towards the total number of layers
+        static_assert(NUM_LAYERS >= 2); // At least input and output layer are required
+        static constexpr TI NUM_HIDDEN_LAYERS = NUM_LAYERS - 2;
+        static constexpr T_TI HIDDEN_DIM = T_HIDDEN_DIM;
+        static constexpr auto HIDDEN_ACTIVATION_FUNCTION = T_HIDDEN_ACTIVATION_FUNCTION;
+        static constexpr auto OUTPUT_ACTIVATION_FUNCTION = T_OUTPUT_ACTIVATION_FUNCTION;
+        static constexpr TI IS_INPUT = T_IS_INPUT;
+        static constexpr TI IS_OUTPUT = T_IS_OUTPUT;
+
+        using LAYER_INITIALIZER = T_LAYER_INITIALIZER;
+
+    };
+    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 constexpr TI INPUT_DIM = get_last(INPUT_SHAPE{});
+        template <typename NEW_INPUT_SHAPE>
+        using OUTPUT_SHAPE_FACTORY = tensor::Replace<NEW_INPUT_SHAPE, CONFIG::OUTPUT_DIM, length(NEW_INPUT_SHAPE{})-1>;
+        using OUTPUT_SHAPE = OUTPUT_SHAPE_FACTORY<INPUT_SHAPE>;
+        static constexpr TI INTERNAL_BATCH_SIZE = get<0>(tensor::CumulativeProduct<tensor::PopBack<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.
+        static constexpr TI NUM_WEIGHTS = CONFIG::OUTPUT_DIM * INPUT_DIM + CONFIG::OUTPUT_DIM;
+
+        using INPUT_LAYER_CONFIG  = nn::layers::dense::Configuration<TYPE_POLICY, TI, CONFIG::HIDDEN_DIM, CONFIG::HIDDEN_ACTIVATION_FUNCTION, typename CONFIG::LAYER_INITIALIZER, utils::typing::conditional_t<CONFIG::IS_INPUT, nn::parameters::groups::Input, nn::parameters::groups::Normal>>;
+        using HIDDEN_LAYER_CONFIG = nn::layers::dense::Configuration<TYPE_POLICY, TI, CONFIG::HIDDEN_DIM, CONFIG::HIDDEN_ACTIVATION_FUNCTION, typename CONFIG::LAYER_INITIALIZER, nn::parameters::groups::Normal>;
+        using OUTPUT_LAYER_CONFIG = nn::layers::dense::Configuration<TYPE_POLICY, TI, CONFIG::OUTPUT_DIM, CONFIG::OUTPUT_ACTIVATION_FUNCTION, typename CONFIG::LAYER_INITIALIZER, utils::typing::conditional_t<CONFIG::IS_OUTPUT ,nn::parameters::groups::Output, nn::parameters::groups::Normal>>;
+        using INPUT_LAYER = nn::layers::dense::Layer<INPUT_LAYER_CONFIG, CAPABILITY, INPUT_SHAPE>;
+        using HIDDEN_LAYER = nn::layers::dense::Layer<HIDDEN_LAYER_CONFIG, CAPABILITY, typename INPUT_LAYER::SPEC::OUTPUT_SHAPE>;
+        using OUTPUT_LAYER = nn::layers::dense::Layer<OUTPUT_LAYER_CONFIG, CAPABILITY, typename HIDDEN_LAYER::SPEC::OUTPUT_SHAPE>;
+    };
+
+    template<typename SPEC_1, typename SPEC_2>
+    constexpr bool check_spec_memory =
+            // utils::typing::is_same_v<typename SPEC_1::T, typename SPEC_2::T>
+            SPEC_1::INPUT_DIM == SPEC_2::INPUT_DIM
+            && SPEC_1::OUTPUT_DIM == SPEC_2::OUTPUT_DIM
+            && SPEC_1::NUM_LAYERS == SPEC_2::NUM_LAYERS
+            && SPEC_1::HIDDEN_DIM == SPEC_2::HIDDEN_DIM;
+    template<typename SPEC_1, typename SPEC_2>
+    constexpr bool check_spec =
+            check_spec_memory<SPEC_1, SPEC_2>
+            && SPEC_1::HIDDEN_ACTIVATION_FUNCTION == SPEC_2::HIDDEN_ACTIVATION_FUNCTION
+            && SPEC_1::OUTPUT_ACTIVATION_FUNCTION == SPEC_2::OUTPUT_ACTIVATION_FUNCTION;
+
+//    template <typename T_CAPABILITY, typename T_SPEC>
+//    struct CapabilitySpecification: T_SPEC, T_CAPABILITY{
+//        using CAPABILITY = T_CAPABILITY;
+//        using PARAMETER_TYPE = typename CAPABILITY::PARAMETER_TYPE;
+//    };
+    struct State{};
+
+    // T_LAYER_PROTOTYPE is any of INPUT_LAYER, HIDDEN_LAYER, or OUTPUT_LAYER. It is assumed that the buffer for all of them are the same (likely empty)
+    template<typename T_SPEC, bool T_DYNAMIC_ALLOCATION>
+    struct NeuralNetworkBuffersSpecification{
+        using SPEC = T_SPEC;
+        using TI = typename SPEC::TI;
+        static constexpr bool DYNAMIC_ALLOCATION = T_DYNAMIC_ALLOCATION;
+        static constexpr TI INTERNAL_BATCH_SIZE = SPEC::INTERNAL_BATCH_SIZE;
+        static constexpr TI DIM = SPEC::HIDDEN_DIM;
+    };
+
+    struct BuffersTag {
+        using CATEGORY = numeric_types::categories::Accumulator;
+    };
+
+    template<typename T_BUFFER_SPEC>
+    struct NeuralNetworkBuffers{
+        using BUFFER_SPEC = T_BUFFER_SPEC;
+        using MLP_SPEC = typename BUFFER_SPEC::SPEC;
+        using TYPE_POLICY = typename MLP_SPEC::TYPE_POLICY;
+        using TI = typename MLP_SPEC::TI;
+        static constexpr TI INTERNAL_BATCH_SIZE = T_BUFFER_SPEC::INTERNAL_BATCH_SIZE;
+        using T = typename MLP_SPEC::TYPE_POLICY::template GET<BuffersTag>;
+        using TICK_TOCK_CONTAINER_SPEC = matrix::Specification<T, TI, INTERNAL_BATCH_SIZE, BUFFER_SPEC::DIM, BUFFER_SPEC::DYNAMIC_ALLOCATION>;
+        using TICK_TOCK_CONTAINER_TYPE = Matrix<TICK_TOCK_CONTAINER_SPEC>;
+        TICK_TOCK_CONTAINER_TYPE tick;
+        TICK_TOCK_CONTAINER_TYPE tock;
+        using LayerBuffer = typename MLP_SPEC::INPUT_LAYER::template Buffer<BUFFER_SPEC::DYNAMIC_ALLOCATION>;
+        LayerBuffer layer_buffer;
+    };
+
+    template<typename T_SPEC>
+    struct NeuralNetworkForward{
+        using SPEC = T_SPEC;
+        using TYPE_POLICY = typename SPEC::TYPE_POLICY;
+        using TI = typename SPEC::TI;
+        // Could be dependent on the capability but in this case the buffer-requirements of forward and backward are the same
+
+        // Convenience
+        static_assert(SPEC::NUM_LAYERS >= 2); // At least input and output layer are required
+        static constexpr TI NUM_HIDDEN_LAYERS = SPEC::NUM_LAYERS - 2;
+        static_assert(SPEC::NUM_HIDDEN_LAYERS == NUM_HIDDEN_LAYERS);
+
+        // Interface
+        // static constexpr TI  INPUT_DIM = SPEC::INPUT_LAYER_SPEC::INPUT_DIM;
+        // static constexpr TI OUTPUT_DIM = SPEC::OUTPUT_LAYER_SPEC::OUTPUT_DIM;
+        static constexpr TI NUM_WEIGHTS = SPEC::INPUT_LAYER::NUM_WEIGHTS + SPEC::HIDDEN_LAYER::NUM_WEIGHTS * NUM_HIDDEN_LAYERS + SPEC::OUTPUT_LAYER::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>;
+
+        // Storage
+        typename SPEC::INPUT_LAYER input_layer;
+        typename SPEC::HIDDEN_LAYER  hidden_layers[NUM_HIDDEN_LAYERS];
+        typename SPEC::OUTPUT_LAYER  output_layer;
+
+        template<bool DYNAMIC_ALLOCATION=true>
+        using State = mlp::State;
+        template<bool DYNAMIC_ALLOCATION=true>
+        using Buffer = NeuralNetworkBuffers<NeuralNetworkBuffersSpecification<SPEC, DYNAMIC_ALLOCATION>>;
+    };
+
+    template<typename SPEC>
+    struct NeuralNetworkBackward: public NeuralNetworkForward<SPEC>{
+//        static constexpr typename SPEC::TI BATCH_SIZE = SPEC::BATCH_SIZE;
+    };
+    template<typename SPEC>
+    struct NeuralNetworkGradient: public NeuralNetworkBackward<SPEC>{};
+
+    template<typename CONFIG, typename CAPABILITY, typename INPUT_SHAPE>
+    using NeuralNetwork =
+        typename utils::typing::conditional_t<CAPABILITY::TAG == nn::LayerCapability::Forward, NeuralNetworkForward<Specification<CONFIG, CAPABILITY, INPUT_SHAPE>>,
+        typename utils::typing::conditional_t<CAPABILITY::TAG == nn::LayerCapability::Backward, NeuralNetworkBackward<Specification<CONFIG, CAPABILITY, INPUT_SHAPE>>,
+        typename utils::typing::conditional_t<CAPABILITY::TAG == nn::LayerCapability::Gradient, NeuralNetworkGradient<Specification<CONFIG, CAPABILITY, INPUT_SHAPE>>, void>>>;
+
+    template <typename T_CONFIG>
+    struct BindConfiguration{
+        template <typename CAPABILITY, typename INPUT_SHAPE>
+        using Layer = nn_models::mlp::NeuralNetwork<T_CONFIG, CAPABILITY, INPUT_SHAPE>;
+    };
+
+    template <typename CAPABILITY, typename MODULE, typename INPUT_SHAPE>
+    struct Build: NeuralNetwork<CAPABILITY, MODULE, INPUT_SHAPE>{
+        template <typename TI, TI BATCH_SIZE>
+        struct CHANGE_BATCH_SIZE_IMPL{
+            using NEW_INPUT_SHAPE = tensor::Replace<INPUT_SHAPE, BATCH_SIZE, 1>;
+            using CHANGE_BATCH_SIZE = Build<CAPABILITY, MODULE, NEW_INPUT_SHAPE>;
+        };
+        template <typename TI, TI BATCH_SIZE>
+        using CHANGE_BATCH_SIZE = typename CHANGE_BATCH_SIZE_IMPL<TI, BATCH_SIZE>::CHANGE_BATCH_SIZE;
+        template <typename NEW_CAPABILITY>
+        using CHANGE_CAPABILITY = Build<NEW_CAPABILITY, MODULE, INPUT_SHAPE>;
+    };
+
+
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+#endif

hyperdrone/_vendor/rl-tools/include/rl_tools/nn_models/mlp/operations_cuda.h

@@ -0,0 +1,10 @@
+#include "../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_NN_MODELS_MLP_OPERATIONS_CUDA_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_NN_MODELS_MLP_OPERATIONS_CUDA_H
+
+#include "operations_generic.h"
+#include "../../devices/cuda.h"
+#include <cuda.h>
+
+#endif

hyperdrone/_vendor/rl-tools/include/rl_tools/nn_models/mlp/operations_dummy.h

@@ -0,0 +1,8 @@
+#include "../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_NN_MODELS_MLP_OPERATIONS_DUMMY_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_NN_MODELS_MLP_OPERATIONS_DUMMY_H
+
+#include "operations_generic.h"
+
+#endif

hyperdrone/_vendor/rl-tools/include/rl_tools/nn_models/mlp/operations_generic.h

@@ -0,0 +1,407 @@
+#include "../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_NN_MODELS_MLP_OPERATIONS_GENERIC_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_NN_MODELS_MLP_OPERATIONS_GENERIC_H
+
+#include "network.h"
+#include "../behaviors/skip_parameter_gradients.h"
+#include "../../nn/operations_generic.h"
+#include "../../nn/parameters/operations_generic.h"
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools {
+    template<typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void malloc(DEVICE& device, nn_models::mlp::NeuralNetworkForward<SPEC>& network) {
+        malloc(device, network.input_layer);
+        for (typename DEVICE::index_t layer_i = 0; layer_i < SPEC::NUM_HIDDEN_LAYERS; layer_i++){
+            malloc(device, network.hidden_layers[layer_i]);
+        }
+        malloc(device, network.output_layer);
+    }
+    template<typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void free(DEVICE& device, nn_models::mlp::NeuralNetworkForward<SPEC>& network) {
+        free(device, network.input_layer);
+        for (typename DEVICE::index_t layer_i = 0; layer_i < SPEC::NUM_HIDDEN_LAYERS; layer_i++){
+            free(device, network.hidden_layers[layer_i]);
+        }
+        free(device, network.output_layer);
+    }
+    template<typename DEVICE>
+    RL_TOOLS_FUNCTION_PLACEMENT void malloc(DEVICE& device, nn_models::mlp::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_models::mlp::State& source, nn_models::mlp::State& target){}
+    template<typename DEVICE, typename SPEC, typename RNG, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT void reset(DEVICE& device, const nn_models::mlp::NeuralNetworkForward<SPEC>& layer, nn_models::mlp::State& state, RNG&, Mode<MODE> mode = Mode<mode::Default<>>{}) { } // no-op
+    template<typename DEVICE>
+    RL_TOOLS_FUNCTION_PLACEMENT void free(DEVICE& device, nn_models::mlp::State& state) { } // no-op
+    template<typename DEVICE, typename BUFFER_SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void malloc(DEVICE& device, nn_models::mlp::NeuralNetworkBuffers<BUFFER_SPEC>& buffers) {
+        malloc(device, buffers.tick);
+        malloc(device, buffers.tock);
+    }
+    template<typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void free(DEVICE& device, nn_models::mlp::NeuralNetworkBuffers<SPEC>& buffers) {
+        free(device, buffers.tick);
+        free(device, buffers.tock);
+    }
+    template<typename DEVICE, typename SPEC, typename RNG>
+    RL_TOOLS_FUNCTION_PLACEMENT void init_weights(DEVICE& device, nn_models::mlp::NeuralNetworkForward<SPEC>& network, RNG& rng) {
+        init_weights(device, network.input_layer, rng);
+        for (typename DEVICE::index_t layer_i = 0; layer_i < SPEC::NUM_HIDDEN_LAYERS; layer_i++){
+            init_weights(device, network.hidden_layers[layer_i], rng);
+        }
+        init_weights(device, network.output_layer, rng);
+    }
+
+    template <typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT constexpr auto output(DEVICE& device, nn_models::mlp::NeuralNetworkForward<SPEC>& m){
+        auto tensor_flat = to_tensor(device, m.output_layer.output);
+        auto tensor = view_memory<typename SPEC::OUTPUT_SHAPE>(device, tensor_flat);
+        return tensor;
+    }
+
+    // evaluate does not set intermediate outputs and hence can also be called from stateless layers, for register efficiency use forward when working with "Backward" compatible layers
+
+    namespace nn_models::mlp{
+        template <typename MODEL_SPEC, typename INPUT_SPEC, typename OUTPUT_SPEC>
+        RL_TOOLS_FUNCTION_PLACEMENT constexpr bool check_input_output_f(){
+            static_assert(INPUT_SPEC::COLS == MODEL_SPEC::INPUT_DIM);
+            static_assert(INPUT_SPEC::ROWS == OUTPUT_SPEC::ROWS);
+            static_assert(OUTPUT_SPEC::COLS == MODEL_SPEC::OUTPUT_DIM);
+            // static_assert(utils::typing::is_same_v<typename MODEL_SPEC::T, typename INPUT_SPEC::T>);
+            // static_assert(utils::typing::is_same_v<typename INPUT_SPEC::T, typename OUTPUT_SPEC::T>);
+            return true;
+        }
+        template <typename MODEL_SPEC, typename INPUT_SPEC, typename OUTPUT_SPEC>
+        constexpr bool check_input_output = check_input_output_f<MODEL_SPEC, INPUT_SPEC, OUTPUT_SPEC>();
+    }
+    template<typename DEVICE, typename MODEL_SPEC, typename INPUT_SPEC, typename OUTPUT_SPEC, typename BUFFER_MODEL_SPEC, typename RNG, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT void evaluate(DEVICE& device, const nn_models::mlp::NeuralNetworkForward<MODEL_SPEC>& network, const Matrix<INPUT_SPEC>& input, Matrix<OUTPUT_SPEC>& output, nn_models::mlp::NeuralNetworkBuffers<BUFFER_MODEL_SPEC>& buffers, RNG& rng, const Mode<MODE>& mode = Mode<mode::Default<>>{}){
+        static_assert(nn_models::mlp::check_input_output<MODEL_SPEC, INPUT_SPEC, OUTPUT_SPEC>);
+        constexpr auto BATCH_SIZE = INPUT_SPEC::ROWS;
+        static_assert(OUTPUT_SPEC::ROWS >= BATCH_SIZE);
+        static_assert(BUFFER_MODEL_SPEC::DIM >= MODEL_SPEC::HIDDEN_DIM);
+        matrix::ViewSpec<BATCH_SIZE, MODEL_SPEC::HIDDEN_DIM> hidden_vs;
+        {
+            auto output_buffer_view = view(device, buffers.tick, matrix::ViewSpec<BATCH_SIZE, MODEL_SPEC::HIDDEN_DIM>{});
+            evaluate(device, network.input_layer, input, output_buffer_view, buffers.layer_buffer, rng, mode);
+        }
+        for (typename DEVICE::index_t layer_i = 0; layer_i < MODEL_SPEC::NUM_HIDDEN_LAYERS; layer_i++){
+            auto& input_buffer = (layer_i % 2 == 0) ? buffers.tick : buffers.tock;
+            auto input_buffer_view = view(device, input_buffer, hidden_vs);
+            auto& output_buffer = (layer_i % 2 == 0) ? buffers.tock : buffers.tick;
+            auto output_buffer_view = view(device, output_buffer, hidden_vs);
+            evaluate(device, network.hidden_layers[layer_i], input_buffer_view, output_buffer_view, buffers.layer_buffer, rng, mode);
+        }
+        if constexpr(MODEL_SPEC::NUM_HIDDEN_LAYERS % 2 == 0){
+            auto input_buffer_view = view(device, buffers.tick, hidden_vs);
+            evaluate(device, network.output_layer, input_buffer_view, output, buffers.layer_buffer, rng, mode);
+        } else {
+            auto input_buffer_view = view(device, buffers.tock, hidden_vs);
+            evaluate(device, network.output_layer, input_buffer_view, output, buffers.layer_buffer, rng, mode);
+        }
+    }
+    template<typename DEVICE, typename MODEL_SPEC, typename INPUT_SPEC, typename BUFFER_SPEC, typename RNG, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT void forward(DEVICE& device, nn_models::mlp::NeuralNetworkGradient<MODEL_SPEC>& network, const Matrix<INPUT_SPEC>& input, nn_models::mlp::NeuralNetworkBuffers<BUFFER_SPEC>& buffers, RNG& rng, const Mode<MODE>& mode = Mode<mode::Default<>>{}) {
+        nn::layers::dense::Buffer layer_buffer;
+        forward(device, network.input_layer, input, layer_buffer, rng, mode);
+
+        auto current_output = network.input_layer.output;
+        for (typename DEVICE::index_t layer_i = 0; layer_i < MODEL_SPEC::NUM_HIDDEN_LAYERS; layer_i++){
+            forward(device, network.hidden_layers[layer_i], current_output, layer_buffer, rng, mode);
+            current_output = network.hidden_layers[layer_i].output;
+        }
+        forward(device, network.output_layer, current_output, layer_buffer, rng, mode);
+    }
+    template<typename DEVICE, typename MODEL_SPEC, typename INPUT_SPEC, typename OUTPUT_SPEC, typename BUFFER_SPEC, typename RNG, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT void forward(DEVICE& device, nn_models::mlp::NeuralNetworkGradient<MODEL_SPEC>& network, const Matrix<INPUT_SPEC>& input, Matrix<OUTPUT_SPEC>& output, nn_models::mlp::NeuralNetworkBuffers<BUFFER_SPEC>& buffers, RNG& rng, const Mode<MODE>& mode = Mode<mode::Default<>>{}) {
+        static_assert(nn_models::mlp::check_input_output<MODEL_SPEC, INPUT_SPEC, OUTPUT_SPEC>);
+        forward(device, network, input, buffers, rng, mode);
+        copy(device, device, network.output_layer.output, output);
+    }
+
+    template<typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void zero_gradient(DEVICE& device, nn_models::mlp::NeuralNetworkGradient<SPEC>& network) {
+        zero_gradient(device, network.input_layer);
+        for(typename DEVICE::index_t i = 0; i < SPEC::NUM_HIDDEN_LAYERS; i++){
+            zero_gradient(device, network.hidden_layers[i]);
+        }
+        zero_gradient(device, network.output_layer);
+    }
+    template<typename DEVICE, typename SOURCE_SPEC, typename TARGET_SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void add_gradient(DEVICE& device, nn_models::mlp::NeuralNetworkGradient<SOURCE_SPEC>& source, nn_models::mlp::NeuralNetworkGradient<TARGET_SPEC>& target) {
+        static_assert(SOURCE_SPEC::NUM_HIDDEN_LAYERS == TARGET_SPEC::NUM_HIDDEN_LAYERS);
+        add_gradient(device, source.input_layer, target.input_layer);
+        for(typename DEVICE::index_t i = 0; i < SOURCE_SPEC::NUM_HIDDEN_LAYERS; i++){
+            add_gradient(device, source.hidden_layers[i], target.hidden_layers[i]);
+        }
+        add_gradient(device, source.output_layer, target.output_layer);
+    }
+    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_models::mlp::NeuralNetworkGradient<SOURCE_SPEC>& source, nn_models::mlp::NeuralNetworkGradient<TARGET_SPEC>& target) {
+        static_assert(SOURCE_SPEC::NUM_HIDDEN_LAYERS == TARGET_SPEC::NUM_HIDDEN_LAYERS);
+        copy_gradient(source_device, target_device, source.input_layer, target.input_layer);
+        for(typename SOURCE_DEVICE::index_t i = 0; i < SOURCE_SPEC::NUM_HIDDEN_LAYERS; i++){
+            copy_gradient(source_device, target_device, source.hidden_layers[i], target.hidden_layers[i]);
+        }
+        copy_gradient(source_device, target_device, source.output_layer, target.output_layer);
+    }
+    template<typename DEVICE, typename MODEL_SPEC, typename D_OUTPUT_SPEC, typename D_INPUT_SPEC, typename BUFFER_MODEL_SPEC, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT void backward_input(DEVICE& device, nn_models::mlp::NeuralNetworkBackward<MODEL_SPEC>& network, Matrix<D_OUTPUT_SPEC>& d_output, Matrix<D_INPUT_SPEC>& d_input, nn_models::mlp::NeuralNetworkBuffers<BUFFER_MODEL_SPEC>& buffer, const Mode<MODE>& mode = Mode<mode::Default<>>{}) {
+        // ATTENTION: this modifies d_output (uses it as a buffer for the d_pre_activations
+        static_assert(nn_models::mlp::check_input_output<MODEL_SPEC, D_INPUT_SPEC, D_OUTPUT_SPEC>);
+        constexpr auto BATCH_SIZE = D_INPUT_SPEC::ROWS;
+        static_assert(BUFFER_MODEL_SPEC::INTERNAL_BATCH_SIZE == BATCH_SIZE);
+        static_assert(BUFFER_MODEL_SPEC::DIM >= MODEL_SPEC::HIDDEN_DIM);
+        using TI = typename DEVICE::index_t;
+
+        backward_input(device, network.output_layer, d_output, buffer.tick, buffer.layer_buffer, mode);
+        for (TI layer_i_plus_one = MODEL_SPEC::NUM_HIDDEN_LAYERS; layer_i_plus_one > 0; layer_i_plus_one--){
+            TI layer_i = layer_i_plus_one - 1;
+            if(layer_i % 2 == (MODEL_SPEC::NUM_HIDDEN_LAYERS - 1) % 2){ // we are starting with the last hidden layer where the result should go to tock
+                backward_input(device, network.hidden_layers[layer_i], buffer.tick, buffer.tock, buffer.layer_buffer, mode);
+            } else {
+                backward_input(device, network.hidden_layers[layer_i], buffer.tock, buffer.tick, buffer.layer_buffer, mode);
+            }
+        }
+        auto& target_d_output_buffer = (MODEL_SPEC::NUM_HIDDEN_LAYERS % 2 == 0) ? buffer.tick : buffer.tock;
+        backward_input(device, network.input_layer, target_d_output_buffer, d_input, buffer.layer_buffer, mode);
+    }
+
+    template<typename DEVICE, typename MODEL_SPEC, typename INPUT_SPEC, typename D_OUTPUT_SPEC, typename BUFFER_MODEL_SPEC, typename D_INPUT_SPEC, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT void backward_full(DEVICE& device, nn_models::mlp::NeuralNetworkGradient<MODEL_SPEC>& network, const Matrix<INPUT_SPEC>& input, Matrix<D_OUTPUT_SPEC>& d_output, Matrix<D_INPUT_SPEC>& d_input, nn_models::mlp::NeuralNetworkBuffers<BUFFER_MODEL_SPEC>& buffer, const Mode<MODE>& mode = Mode<mode::Default<>>{}) {
+        // ATTENTION: this modifies d_output (uses it as a buffer for the d_pre_activations
+        static_assert(nn_models::mlp::check_input_output<MODEL_SPEC, D_INPUT_SPEC, D_OUTPUT_SPEC>);
+        static_assert(nn_models::mlp::check_input_output<MODEL_SPEC, INPUT_SPEC, D_OUTPUT_SPEC>);
+        constexpr auto BATCH_SIZE = D_INPUT_SPEC::ROWS;
+        static_assert(BUFFER_MODEL_SPEC::INTERNAL_BATCH_SIZE == BATCH_SIZE);
+        static_assert(BUFFER_MODEL_SPEC::DIM >= MODEL_SPEC::HIDDEN_DIM);
+
+        if constexpr(!nn_models::backpropagation::ParameterGradients<MODEL_SPEC, MODE>::VALUE){
+            backward_input(device, network, d_output, d_input, buffer, mode);
+        }
+        else{
+            {
+                auto& previous_output = MODEL_SPEC::NUM_HIDDEN_LAYERS > 0 ? network.hidden_layers[MODEL_SPEC::NUM_HIDDEN_LAYERS - 1].output : network.input_layer.output;
+                backward_full(device, network.output_layer, previous_output, d_output, buffer.tick, buffer.layer_buffer, mode);
+            }
+            for (typename DEVICE::index_t layer_i_plus_one = MODEL_SPEC::NUM_HIDDEN_LAYERS; layer_i_plus_one > 0; layer_i_plus_one--){
+                typename DEVICE::index_t layer_i = layer_i_plus_one - 1;
+                auto& previous_output_hidden = layer_i > 0 ? network.hidden_layers[layer_i - 1].output : network.input_layer.output;
+                if(layer_i % 2 == (MODEL_SPEC::NUM_HIDDEN_LAYERS - 1) % 2){ // we are starting with the last hidden layer where the result should go to tock
+                    backward_full(device, network.hidden_layers[layer_i], previous_output_hidden, buffer.tick, buffer.tock, buffer.layer_buffer, mode);
+                } else {
+                    backward_full(device, network.hidden_layers[layer_i], previous_output_hidden, buffer.tock, buffer.tick, buffer.layer_buffer, mode);
+                }
+            }
+            if constexpr(MODEL_SPEC::NUM_HIDDEN_LAYERS % 2 == 0){
+                backward_full(device, network.input_layer, input, buffer.tick, d_input, buffer.layer_buffer, mode);
+            } else {
+                backward_full(device, network.input_layer, input, buffer.tock, d_input, buffer.layer_buffer, mode);
+            }
+        }
+    }
+    template<typename DEVICE, typename MODEL_SPEC, typename INPUT_SPEC, typename D_OUTPUT_SPEC, typename BUFFER_MODEL_SPEC, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT void backward(DEVICE& device, nn_models::mlp::NeuralNetworkGradient<MODEL_SPEC>& network, const Matrix<INPUT_SPEC>& input, Matrix<D_OUTPUT_SPEC>& d_output, nn_models::mlp::NeuralNetworkBuffers<BUFFER_MODEL_SPEC>& buffer, const Mode<MODE>& mode = Mode<mode::Default<>>{}) {
+        // ATTENTION: this modifies d_output (uses it as a buffer for the d_pre_activations
+        static_assert(nn_models::mlp::check_input_output<MODEL_SPEC, INPUT_SPEC, D_OUTPUT_SPEC>);
+        constexpr auto BATCH_SIZE = D_OUTPUT_SPEC::ROWS;
+        static_assert(BUFFER_MODEL_SPEC::INTERNAL_BATCH_SIZE == BATCH_SIZE);
+        static_assert(BUFFER_MODEL_SPEC::DIM >= MODEL_SPEC::HIDDEN_DIM);
+
+        if constexpr(nn_models::backpropagation::ParameterGradients<MODEL_SPEC, MODE>::VALUE){
+            auto previous_output = MODEL_SPEC::NUM_HIDDEN_LAYERS > 0 ? network.hidden_layers[MODEL_SPEC::NUM_HIDDEN_LAYERS - 1].output : network.input_layer.output;
+            backward_full(device, network.output_layer, previous_output, d_output, buffer.tick, buffer.layer_buffer, mode);
+            for (typename DEVICE::index_t layer_i_plus_one = MODEL_SPEC::NUM_HIDDEN_LAYERS; layer_i_plus_one > 0; layer_i_plus_one--){
+                typename DEVICE::index_t layer_i = layer_i_plus_one - 1;
+                previous_output = layer_i > 0 ? network.hidden_layers[layer_i - 1].output : network.input_layer.output;
+                if(layer_i % 2 == (MODEL_SPEC::NUM_HIDDEN_LAYERS - 1) % 2){ // we are starting with the last hidden layer where the result should go to tock
+                    backward_full(device, network.hidden_layers[layer_i], previous_output, buffer.tick, buffer.tock, buffer.layer_buffer, mode);
+                } else {
+                    backward_full(device, network.hidden_layers[layer_i], previous_output, buffer.tock, buffer.tick, buffer.layer_buffer, mode);
+                }
+            }
+            if constexpr(MODEL_SPEC::NUM_HIDDEN_LAYERS % 2 == 0){
+                backward(device, network.input_layer, input, buffer.tick, buffer.layer_buffer, mode);
+            } else {
+                backward(device, network.input_layer, input, buffer.tock, buffer.layer_buffer, mode);
+            }
+        }
+    }
+//    template<typename DEVICE, typename MODEL_SPEC, typename INPUT_SPEC, typename D_OUTPUT_SPEC, typename BUFFER_MODEL_SPEC>
+//    void backward_param(DEVICE& device, nn_models::mlp::NeuralNetworkGradient<MODEL_SPEC>& network, const Matrix<INPUT_SPEC>& input, Matrix<D_OUTPUT_SPEC>& d_output, nn_models::mlp::NeuralNetworkBuffers<BUFFER_MODEL_SPEC> buffer) {
+//        backward(device, network, input, d_output, buffer);
+//    }
+
+    template<typename DEVICE, typename SPEC, typename OPTIMIZER>
+    RL_TOOLS_FUNCTION_PLACEMENT void update(DEVICE& device, nn_models::mlp::NeuralNetworkGradient<SPEC>& network, OPTIMIZER& optimizer) {
+        update(device, network.input_layer, optimizer);
+        for(typename DEVICE::index_t layer_i = 0; layer_i < SPEC::NUM_HIDDEN_LAYERS; layer_i++){
+            update(device, network.hidden_layers[layer_i], optimizer);
+        }
+        update(device, network.output_layer, optimizer);
+    }
+
+    template<typename DEVICE, typename SPEC, typename OPTIMIZER>
+    RL_TOOLS_FUNCTION_PLACEMENT void _reset_optimizer_state(DEVICE& device, nn_models::mlp::NeuralNetworkGradient<SPEC>& network, OPTIMIZER& optimizer) {
+        // this function is marked with a underscore because it should usually be called from the reset_optimizer_state function of the optimizer to have one coherent entrypoint for resetting the optimizer state in the optimizer and in the model
+        _reset_optimizer_state(device, network.input_layer, optimizer);
+        for(typename DEVICE::index_t layer_i = 0; layer_i < SPEC::NUM_HIDDEN_LAYERS; layer_i++){
+            _reset_optimizer_state(device, network.hidden_layers[layer_i], optimizer);
+        }
+        _reset_optimizer_state(device, network.output_layer, optimizer);
+    }
+
+    // The following copy operators are more powerful than the default copy assignment operator in that they can e.g. copy between networks with different activation functions
+    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_models::mlp::NeuralNetworkForward<SOURCE_SPEC>& source, nn_models::mlp::NeuralNetworkForward<TARGET_SPEC>& target){
+        static_assert(rl_tools::nn_models::mlp::check_spec_memory<SOURCE_SPEC, TARGET_SPEC>, "The source and target network must have the same structure");
+        copy(source_device, target_device, source.input_layer, target.input_layer);
+        for(typename SOURCE_SPEC::TI layer_i = 0; layer_i <  SOURCE_SPEC::NUM_HIDDEN_LAYERS; layer_i++){
+            copy(source_device, target_device, source.hidden_layers[layer_i], target.hidden_layers[layer_i]);
+        }
+        copy(source_device, target_device, source.output_layer, target.output_layer);
+    }
+    template<typename SOURCE_DEVICE, typename TARGET_DEVICE,  typename SOURCE, typename TARGET_SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void copy_from_generic(SOURCE_DEVICE& source_device, TARGET_DEVICE& target_device, const SOURCE& source, nn_models::mlp::NeuralNetworkForward<TARGET_SPEC>& target){
+        static_assert(rl_tools::nn_models::mlp::check_spec_memory<SOURCE, TARGET_SPEC>, "The source and target network must have the same structure");
+        copy(source_device, target_device, source.input_layer, target.input_layer);
+        for(typename SOURCE::TI layer_i = 0; layer_i <  SOURCE::NUM_HIDDEN_LAYERS; layer_i++){
+            copy(source_device, target_device, source.hidden_layers[layer_i], target.hidden_layers[layer_i]);
+        }
+        copy(source_device, target_device, source.output_layer, target.output_layer);
+    }
+
+//    template<typename SOURCE_DEVICE, typename TARGET_DEVICE, typename SOURCE_SPEC, typename TARGET_SPEC>
+//    void copy(SOURCE_DEVICE& source_device, TARGET_DEVICE& target_device, const  nn_models::mlp::NeuralNetworkAdam<SOURCE_SPEC>& source, nn_models::mlp::NeuralNetworkAdam<TARGET_SPEC>& target){
+//        static_assert(rl_tools::nn_models::mlp::check_spec_memory<SOURCE_SPEC, TARGET_SPEC>, "The source and target network must have the same structure");
+//        copy(source_device, target_device, (nn_models::mlp::NeuralNetworkForward<SOURCE_SPEC>&)source, (nn_models::mlp::NeuralNetworkForward<TARGET_SPEC>&)target);
+//    }
+
+    template<typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void reset_forward_state(DEVICE& device, nn_models::mlp::NeuralNetworkForward<SPEC>& n){
+        reset_forward_state(device, n.input_layer);
+        for(typename DEVICE::index_t layer_i = 0; layer_i <  SPEC::NUM_HIDDEN_LAYERS; layer_i++){
+            reset_forward_state(device, n.hidden_layers[layer_i]);
+        }
+        reset_forward_state(device, n.output_layer);
+    }
+
+    template<typename DEVICE, typename SPEC_1, typename SPEC_2>
+    RL_TOOLS_FUNCTION_PLACEMENT typename SPEC_1::TYPE_POLICY::DEFAULT abs_diff(DEVICE& device, nn_models::mlp::NeuralNetworkForward<SPEC_1>& n1, const nn_models::mlp::NeuralNetworkForward<SPEC_2>& n2){
+        using T = typename SPEC_1::TYPE_POLICY::DEFAULT;
+        static_assert(rl_tools::nn_models::mlp::check_spec_memory<SPEC_1, SPEC_2>, "The source and target network must have the same structure");
+        T acc = 0;
+
+        acc += abs_diff(device, n1.output_layer, n2.output_layer);
+        for(typename DEVICE::index_t layer_i = 0; layer_i < SPEC_1::NUM_HIDDEN_LAYERS; layer_i++){
+            acc += abs_diff(device, n1.hidden_layers[layer_i], n2.hidden_layers[layer_i]);
+        }
+        acc += abs_diff(device, n1.input_layer, n2.input_layer);
+        return acc;
+    }
+    template <typename DEVICE, typename SPEC, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT bool is_nan(DEVICE& device, const rl_tools::nn_models::mlp::NeuralNetworkForward<SPEC>& n, const Mode<MODE>& mode = Mode<mode::Default<>>{}){
+        bool found_nan = false;
+        found_nan = found_nan || is_nan(device, n.input_layer, mode);
+        for(typename DEVICE::index_t layer_i = 0; layer_i < SPEC::NUM_HIDDEN_LAYERS; layer_i++){
+            found_nan = found_nan || is_nan(device, n.hidden_layers[layer_i], mode);
+        }
+        found_nan = found_nan || is_nan(device, n.output_layer, mode);
+        return found_nan;
+    }
+    template<typename DEVICE>
+    RL_TOOLS_FUNCTION_PLACEMENT auto abs_diff(DEVICE& device, nn_models::mlp::State& s1, nn_models::mlp::State& s2){
+        return 0;
+    }
+    template<typename DEVICE, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT bool is_nan(DEVICE& device, nn_models::mlp::State& state, const Mode<MODE>& mode = Mode<mode::Default<>>{}){
+        return false;
+    }
+    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_models::mlp::NeuralNetworkBuffers<SOURCE_SPEC>& source, nn_models::mlp::NeuralNetworkBuffers<TARGET_SPEC>& target){
+        copy(source_device, target_device, source.tick, target.tick);
+        copy(source_device, target_device, source.tock, target.tock);
+    }
+    template<typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT constexpr auto& output(nn_models::mlp::NeuralNetworkGradient<SPEC>& nn){
+        return nn.output_layer.output;
+    }
+    template<typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT constexpr auto output(DEVICE& device, nn_models::mlp::NeuralNetworkGradient<SPEC>& nn){
+        return output(device, nn.output_layer);
+    }
+    template <typename DEVICE, typename BUFFER_SPEC, typename RNG>
+    RL_TOOLS_FUNCTION_PLACEMENT void sample(DEVICE& device, nn_models::mlp::NeuralNetworkBuffers<BUFFER_SPEC>& buffers, RNG& rng){ }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+
+// Tensor proxies
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools{
+    template<typename DEVICE, typename SPEC, typename INPUT_SPEC, typename OUTPUT_SPEC, typename BUFFER_MODEL_SPEC, typename RNG, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT void evaluate(DEVICE& device, const nn_models::mlp::NeuralNetworkForward<SPEC>& model, const Tensor<INPUT_SPEC>& input, Tensor<OUTPUT_SPEC>& output,nn_models::mlp::NeuralNetworkBuffers<BUFFER_MODEL_SPEC>& 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, model, matrix_view_input, matrix_view_output, buffer, rng, mode);
+    }
+    template<typename DEVICE, typename SPEC, typename INPUT_SPEC, typename OUTPUT_SPEC, typename BUFFER_MODEL_SPEC, typename RNG, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT void evaluate_step(DEVICE& device, const nn_models::mlp::NeuralNetworkForward<SPEC>& model, const Tensor<INPUT_SPEC>& input, nn_models::mlp::State, Tensor<OUTPUT_SPEC>& output,nn_models::mlp::NeuralNetworkBuffers<BUFFER_MODEL_SPEC>& 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, model, matrix_view_input, matrix_view_output, buffer, rng, mode);
+    }
+    template<typename DEVICE, typename SPEC, typename INPUT_SPEC, typename OUTPUT_SPEC, typename BUFFER_MODEL_SPEC, typename RNG, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT void forward(DEVICE& device, nn_models::mlp::NeuralNetworkBackward<SPEC>& model, const Tensor<INPUT_SPEC>& input, Tensor<OUTPUT_SPEC>& output,nn_models::mlp::NeuralNetworkBuffers<BUFFER_MODEL_SPEC>& 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, model, matrix_view_input, matrix_view_output, buffer, rng, mode);
+    }
+    template<typename DEVICE, typename SPEC, typename INPUT_SPEC, typename BUFFER_MODEL_SPEC, typename RNG, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT void forward(DEVICE& device, nn_models::mlp::NeuralNetworkGradient<SPEC>& model, const Tensor<INPUT_SPEC>& input,nn_models::mlp::NeuralNetworkBuffers<BUFFER_MODEL_SPEC>& buffer, RNG& rng, const Mode<MODE>& mode = Mode<mode::Default<>>{}) {
+        auto matrix_view_input = matrix_view(device, input);
+        forward(device, model, matrix_view_input, buffer, rng, mode);
+    }
+    template<typename DEVICE, typename SPEC, typename INPUT_SPEC, typename OUTPUT_SPEC, typename BUFFER_MODEL_SPEC, typename RNG, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT void forward(DEVICE& device, nn_models::mlp::NeuralNetworkGradient<SPEC>& model, const Tensor<INPUT_SPEC>& input, Tensor<OUTPUT_SPEC>& output,nn_models::mlp::NeuralNetworkBuffers<BUFFER_MODEL_SPEC>& 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, model, matrix_view_input, matrix_view_output, buffer, rng, mode);
+    }
+    template<typename DEVICE, typename SPEC, typename D_OUTPUT_SPEC, typename D_INPUT_SPEC, typename BUFFER_MODEL_SPEC, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT void backward_input(DEVICE& device, nn_models::mlp::NeuralNetworkBackward<SPEC>& model, Tensor<D_OUTPUT_SPEC>& d_output, Tensor<D_INPUT_SPEC>& d_input,nn_models::mlp::NeuralNetworkBuffers<BUFFER_MODEL_SPEC>& 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, model, matrix_view_d_output, matrix_view_d_input, buffer, mode);
+    }
+
+    template<typename DEVICE, typename SPEC, typename INPUT_SPEC, typename D_OUTPUT_SPEC, typename BUFFER_MODEL_SPEC, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT void backward(DEVICE& device, nn_models::mlp::NeuralNetworkGradient<SPEC>& model, const Tensor<INPUT_SPEC>& input, Tensor<D_OUTPUT_SPEC>& d_output,nn_models::mlp::NeuralNetworkBuffers<BUFFER_MODEL_SPEC>& 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, model, matrix_view_input, matrix_view_d_output, buffer, mode);
+    }
+
+    template<typename DEVICE, typename SPEC, typename INPUT_SPEC, typename D_OUTPUT_SPEC, typename D_INPUT_SPEC, typename BUFFER_MODEL_SPEC, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT void backward_full(DEVICE& device, nn_models::mlp::NeuralNetworkGradient<SPEC>& model, const Tensor<INPUT_SPEC>& input, Tensor<D_OUTPUT_SPEC>& d_output, Tensor<D_INPUT_SPEC>& d_input,nn_models::mlp::NeuralNetworkBuffers<BUFFER_MODEL_SPEC>& 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, model, 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_models::mlp::NeuralNetworkForward<SPEC>& model){
+        using TI = typename DEVICE::index_t;
+        using T = typename SPEC::TYPE_POLICY::DEFAULT;
+        T return_value = 0;
+        return_value += gradient_norm(device, model.input_layer);
+        for(TI layer_i = 0; layer_i < SPEC::NUM_HIDDEN_LAYERS; layer_i++){
+            return_value += gradient_norm(device, model.hidden_layers[layer_i]);
+        }
+        return_value += gradient_norm(device, model.output_layer);
+        return return_value;
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+
+#endif