hyperverse 0.0.1__py3-none-any.whl

hyperdrone/_vendor/rl-tools/include/rl_tools/nn/layers/dense/operations_esp32/opt.h

@@ -0,0 +1,167 @@
+#include "../../../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_NN_LAYERS_DENSE_OPERATIONS_ESP32_OPT_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_NN_LAYERS_DENSE_OPERATIONS_ESP32_OPT_H
+
+#include "../../../../devices/esp32.h"
+#include "../../../../nn/layers/dense/layer.h"
+#include "../../../../nn/mode.h"
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools{
+    template<typename DEV_SPEC, typename LAYER_SPEC, typename INPUT_SPEC, typename OUTPUT_SPEC, typename RNG, typename MODE = mode::Default<>>
+    void evaluate(devices::esp32::OPT<DEV_SPEC>& device, const nn::layers::dense::LayerForward<LAYER_SPEC>& layer, const Matrix<INPUT_SPEC>& input, Matrix<OUTPUT_SPEC>& output, nn::layers::dense::Buffer&, RNG& rng, const Mode<MODE>& mode = Mode<mode::Default<>>{}) {
+        // For performance reasons: restricted to dense row-major matrices (row-pitch is allowed)
+        static_assert(nn::layers::dense::check_input_output<LAYER_SPEC, INPUT_SPEC, OUTPUT_SPEC>);
+        static_assert(INPUT_SPEC::COL_PITCH == 1);
+        static_assert(OUTPUT_SPEC::COL_PITCH == 1);
+        static_assert(decltype(layer.weights.parameters)::COL_PITCH == 1);
+        static_assert(decltype(layer.biases.parameters)::COL_PITCH == 1);
+//        static_assert(utils::typing::is_same_v<typename LAYER_SPEC::T, float>);
+
+        // Warning do not use the same buffer for input and output!
+        constexpr auto BATCH_SIZE = INPUT_SPEC::ROWS;
+        // static_assert(BATCH_SIZE == 1);
+        using DEVICE = devices::esp32::OPT<DEV_SPEC>;
+        using T = typename LAYER_SPEC::T;
+        {
+
+            T *weights_row;
+            T *input_row = input._data;
+            T *output_row = output._data;
+
+            T *weights_element, *biases_element, *input_element, *output_element;
+
+            T acc;
+            uint32_t weights_row_i, batch_i = BATCH_SIZE, input_i;
+
+            {
+                do{
+                    output_element = output_row;
+                    biases_element = layer.biases.parameters._data;
+
+                    weights_row_i = LAYER_SPEC::OUTPUT_DIM;
+                    weights_row = layer.weights.parameters._data;
+
+                    do{
+                        acc = 0.0f;
+                        input_element = input_row;
+                        weights_element = weights_row;
+
+                        // reduction
+                        input_i = ((uint32_t)LAYER_SPEC::INPUT_DIM) >> 2U;
+                        while (input_i > 0U){
+                            acc += *weights_element++ * *input_element++;
+                            acc += *weights_element++ * *input_element++;
+                            acc += *weights_element++ * *input_element++;
+                            acc += *weights_element++ * *input_element++;
+
+                            input_i--;
+                        }
+                        input_i = ((uint32_t)LAYER_SPEC::INPUT_DIM) % 0x4U;
+                        while (input_i > 0U){
+                            acc += *weights_element++ * *input_element++;
+                            input_i--;
+                        }
+
+                        acc += *biases_element++;
+                        acc = activation<typename DEVICE::SPEC::MATH, T, LAYER_SPEC::ACTIVATION_FUNCTION>(acc);
+
+                        *output_element++ = acc;
+
+                        weights_row_i--;
+
+                        weights_row += decltype(layer.weights.parameters)::ROW_PITCH;
+
+                    }while (weights_row_i > 0U);
+
+                    output_row += OUTPUT_SPEC::ROW_PITCH;
+                    input_row += INPUT_SPEC::ROW_PITCH;
+                    batch_i--;
+                }while (batch_i > 0U);
+            }
+        }
+    }
+
+    template<typename DEV_SPEC, typename LAYER_SPEC, typename INPUT_SPEC, typename OUTPUT_SPEC, typename RNG, typename MODE = mode::Default<>>
+    void forward(devices::esp32::OPT<DEV_SPEC>& device, nn::layers::dense::LayerBackward<LAYER_SPEC>& layer, const Matrix<INPUT_SPEC>& input, Matrix<OUTPUT_SPEC>& output, nn::layers::dense::Buffer&, RNG& rng, const Mode<MODE>& mode = Mode<mode::Default<>>{}) {
+        // For performance reasons: restricted to dense row-major matrices (row-pitch is allowed)
+        static_assert(nn::layers::dense::check_input_output<LAYER_SPEC, INPUT_SPEC, OUTPUT_SPEC>);
+        static_assert(INPUT_SPEC::COL_PITCH == 1);
+        static_assert(OUTPUT_SPEC::COL_PITCH == 1);
+        static_assert(decltype(layer.weights.parameters)::COL_PITCH == 1);
+        static_assert(decltype(layer.biases.parameters)::COL_PITCH == 1);
+//        static_assert(utils::typing::is_same_v<typename LAYER_SPEC::T, float>);
+
+        // Warning do not use the same buffer for input and output!
+        constexpr auto BATCH_SIZE = INPUT_SPEC::ROWS;
+        // static_assert(BATCH_SIZE == 1);
+        using DEVICE = devices::esp32::OPT<DEV_SPEC>;
+        using T = typename LAYER_SPEC::T;
+        using TI = typename DEVICE::index_t;
+        {
+
+            T *weights_row;
+            T *input_row = input._data;
+            T *pre_activations_row = layer.pre_activations._data;
+            T *output_row = output._data;
+
+            T *weights_element, *biases_element, *input_element, *pre_activations_element, *output_element;
+
+            T acc;
+            TI weights_row_i, batch_i = BATCH_SIZE, input_i;
+
+            {
+                do{
+                    pre_activations_element = pre_activations_row;
+                    output_element = output_row;
+                    biases_element = layer.biases.parameters._data;
+
+                    weights_row_i = LAYER_SPEC::OUTPUT_DIM;
+                    weights_row = layer.weights.parameters._data;
+
+                    do{
+                        acc = 0.0f;
+                        input_element = input_row;
+                        weights_element = weights_row;
+
+                        // reduction
+                        input_i = ((TI)LAYER_SPEC::INPUT_DIM) >> 2U;
+                        while (input_i > 0U){
+                            acc += *weights_element++ * *input_element++;
+                            acc += *weights_element++ * *input_element++;
+                            acc += *weights_element++ * *input_element++;
+                            acc += *weights_element++ * *input_element++;
+
+                            input_i--;
+                        }
+                        input_i = ((TI)LAYER_SPEC::INPUT_DIM) % 0x4U;
+                        while (input_i > 0U){
+                            acc += *weights_element++ * *input_element++;
+                            input_i--;
+                        }
+
+                        acc += *biases_element++;
+                        *pre_activations_element++ = acc;
+                        acc = activation<typename DEVICE::SPEC::MATH, T, LAYER_SPEC::ACTIVATION_FUNCTION>(acc);
+
+                        *output_element++ = acc;
+
+                        weights_row_i--;
+
+                        weights_row += decltype(layer.weights.parameters)::ROW_PITCH;
+
+                    }while (weights_row_i > 0U);
+
+                    pre_activations_row += OUTPUT_SPEC::ROW_PITCH;
+                    output_row += OUTPUT_SPEC::ROW_PITCH;
+                    input_row += INPUT_SPEC::ROW_PITCH;
+                    batch_i--;
+                }while (batch_i > 0U);
+            }
+        }
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+
+#endif

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

@@ -0,0 +1,429 @@
+#include "../../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_NN_LAYERS_DENSE_OPERATIONS_GENERIC_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_NN_LAYERS_DENSE_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::dense::LayerForward<SPEC>& layer) {
+        malloc(device, layer.weights);
+        malloc(device, layer.biases);
+    }
+    template<typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void free(DEVICE& device, nn::layers::dense::LayerForward<SPEC>& layer) {
+        free(device, layer.weights);
+        free(device, layer.biases);
+    }
+    template<typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void malloc(DEVICE& device, nn::layers::dense::LayerBackward<SPEC>& layer) {
+        malloc(device, (nn::layers::dense::LayerForward<SPEC>&) layer);
+        malloc(device, layer.pre_activations);
+    }
+    template<typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void free(DEVICE& device, nn::layers::dense::LayerBackward<SPEC>& layer) {
+        free(device, (nn::layers::dense::LayerForward<SPEC>&) layer);
+        free(device, layer.pre_activations);
+    }
+    template<typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void malloc(DEVICE& device, nn::layers::dense::LayerGradient<SPEC>& layer) {
+        malloc(device, (nn::layers::dense::LayerBackward<SPEC>&) layer);
+        malloc(device, layer.output);
+    }
+    template<typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void free(DEVICE& device, nn::layers::dense::LayerGradient<SPEC>& layer) {
+        free(device, (nn::layers::dense::LayerBackward<SPEC>&) layer);
+        free(device, layer.output);
+    }
+    template<typename DEVICE>
+    RL_TOOLS_FUNCTION_PLACEMENT void malloc(DEVICE& device, nn::layers::dense::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::dense::State& source, nn::layers::dense::State& target){}
+    template<typename SPEC, typename DEVICE, typename RNG, typename MODE>
+    RL_TOOLS_FUNCTION_PLACEMENT void reset(DEVICE& device, const nn::layers::dense::LayerForward<SPEC>& layer, nn::layers::dense::State& state, RNG&, Mode<MODE> mode = Mode<mode::Default<>>{}) { } // no-op
+    template<typename DEVICE>
+    RL_TOOLS_FUNCTION_PLACEMENT void free(DEVICE& device, nn::layers::dense::State& state) { } // no-op
+    template<typename DEVICE>
+    RL_TOOLS_FUNCTION_PLACEMENT void malloc(DEVICE& device, nn::layers::dense::Buffer& buffer) { } // no-op
+    template<typename DEVICE>
+    RL_TOOLS_FUNCTION_PLACEMENT void free(DEVICE& device, nn::layers::dense::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::dense::LayerForward<SPEC>& layer, const nn::layers::dense::KaimingUniform<INITIALIZER_SPEC>& initializer, RNG& rng){
+        using T = typename SPEC::TYPE_POLICY::DEFAULT;
+        using TI = typename SPEC::TI;
+        T gain;
+        if constexpr(INITIALIZER_SPEC::INIT_LEGACY){
+            T negative_slope = math::sqrt(device.math, (T)5);
+            gain = math::sqrt(device.math, (T)2.0 / (1 + negative_slope * negative_slope));
+        }
+        else{
+            gain = math::sqrt(device.math, (T)2.0) * INITIALIZER_SPEC::SCALE;
+        }
+        T fan = SPEC::INPUT_DIM;
+        T std = gain / math::sqrt(device.math, fan);
+        T weight_bound = math::sqrt(device.math, (T)3.0) * std;
+        T bias_bound = 1/math::sqrt(device.math, (T)SPEC::INPUT_DIM);
+        using PARAMETER_TYPE = typename decltype(layer.weights.parameters)::SPEC::T;
+        for(TI i = 0; i < SPEC::OUTPUT_DIM; i++) {
+            if constexpr(INITIALIZER_SPEC::INIT_LEGACY) {
+                set(device, layer.biases.parameters, (PARAMETER_TYPE)random::uniform_real_distribution(device.random, -bias_bound, bias_bound, rng), i);
+            }
+            else{
+                set(device, layer.biases.parameters, 0, i);
+            }
+            for(TI j = 0; j < SPEC::INPUT_DIM; j++) {
+                set(device, layer.weights.parameters, (PARAMETER_TYPE)random::uniform_real_distribution(device.random, -weight_bound, weight_bound, rng), i, j);
+            }
+        }
+    }
+    template<typename DEVICE, typename SPEC, typename RNG>
+    RL_TOOLS_FUNCTION_PLACEMENT void init_weights(DEVICE& device, nn::layers::dense::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::dense::LayerForward<LAYER_SPEC>& layer, const Matrix<INPUT_SPEC>& input, Matrix<OUTPUT_SPEC>& output, nn::layers::dense::Buffer&, RNG& rng, const Mode<MODE>& mode = Mode<mode::Default<>>{}) {
+        static_assert(nn::layers::dense::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;
+        using T = typename OUTPUT_SPEC::T;
+        using ACCUMULATOR_TYPE = typename LAYER_SPEC::TYPE_POLICY::template GET<numeric_types::categories::Accumulator>;
+        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++) {
+                ACCUMULATOR_TYPE acc = (ACCUMULATOR_TYPE)get(device, layer.biases.parameters, output_i);
+                for(TI input_i = 0; input_i < LAYER_SPEC::INPUT_DIM; input_i++) {
+                    acc += (ACCUMULATOR_TYPE)get(device, layer.weights.parameters, output_i, input_i) * (ACCUMULATOR_TYPE)get(input, batch_i, input_i);
+                }
+                set(output, batch_i, output_i, (T)activation<typename DEVICE::SPEC::MATH, ACCUMULATOR_TYPE, LAYER_SPEC::ACTIVATION_FUNCTION>(acc));
+            }
+        }
+    }
+
+    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::dense::LayerBackward<LAYER_SPEC>& layer, const Matrix<INPUT_SPEC>& input, Matrix<OUTPUT_SPEC>& output, nn::layers::dense::Buffer& buffer, RNG& rng, const Mode<MODE>& mode = Mode<mode::Default<>>{}){
+        // Warning do not use the same buffer for input and output!
+        static_assert(nn::layers::dense::check_input_output<LAYER_SPEC, INPUT_SPEC, OUTPUT_SPEC>);
+        constexpr auto BATCH_SIZE = INPUT_SPEC::ROWS;
+        using T = typename OUTPUT_SPEC::T;
+        using ACCUMULATOR_TYPE = typename LAYER_SPEC::TYPE_POLICY::template GET<numeric_types::categories::Accumulator>;
+        using TI = typename DEVICE::index_t;
+
+        for(TI batch_i=0; batch_i < BATCH_SIZE; batch_i++){
+            for(TI i = 0; i < LAYER_SPEC::OUTPUT_DIM; i++) {
+                ACCUMULATOR_TYPE acc = (ACCUMULATOR_TYPE)get(device, layer.biases.parameters, i);
+                for(TI j = 0; j < LAYER_SPEC::INPUT_DIM; j++) {
+                    acc += (ACCUMULATOR_TYPE)get(device, layer.weights.parameters, i, j) * (ACCUMULATOR_TYPE)get(input, batch_i, j);
+                }
+                set(layer.pre_activations, batch_i, i, (T)acc);
+                set(output, batch_i, i, (T)activation<typename DEVICE::SPEC::MATH, ACCUMULATOR_TYPE, LAYER_SPEC::ACTIVATION_FUNCTION>(acc));
+            }
+        }
+    }
+#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::dense::LayerGradient<LAYER_SPEC>& layer, const Matrix<INPUT_SPEC>& input, nn::layers::dense::Buffer& buffer, RNG& rng, const Mode<MODE>& mode = Mode<mode::Default<>>{}) {
+        static_assert(nn::layers::dense::check_input_output<LAYER_SPEC, INPUT_SPEC, typename decltype(layer.output)::SPEC>);
+        forward(device, static_cast<nn::layers::dense::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::dense::LayerGradient<LAYER_SPEC>& layer, const Matrix<INPUT_SPEC>& input, Matrix<OUTPUT_SPEC>& output, nn::layers::dense::Buffer& buffer, RNG& rng, const Mode<MODE>& mode = Mode<mode::Default<>>{}) {
+        static_assert(nn::layers::dense::check_input_output<LAYER_SPEC, INPUT_SPEC, OUTPUT_SPEC>);
+        // compile time warning if used
+        forward(device, layer, input, buffer, rng, mode);
+        copy(device, device, layer.output, output);
+    }
+
+#ifndef RL_TOOLS_NN_DISABLE_GENERIC_FORWARD_BACKWARD
+    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::dense::LayerBackward<LAYER_SPEC>& layer, const Matrix<D_OUTPUT_SPEC>& d_output, Matrix<D_INPUT_SPEC>& d_input, nn::layers::dense::Buffer&, const Mode<MODE>& mode = Mode<mode::Default<>>{}){
+        static_assert(nn::layers::dense::check_input_output<LAYER_SPEC, D_INPUT_SPEC, D_OUTPUT_SPEC>);
+        // todo: create sparate function that does not set d_input (to save cost on backward pass for the first layer)
+        using SPEC = LAYER_SPEC;
+        constexpr auto BATCH_SIZE = D_OUTPUT_SPEC::ROWS;
+        using ACCUMULATOR_TYPE = typename SPEC::TYPE_POLICY::template GET<numeric_types::categories::Accumulator>;
+        using TI = typename DEVICE::index_t;
+        for(TI batch_i=0; batch_i < BATCH_SIZE; batch_i++){
+            for(TI output_i = 0; output_i < SPEC::OUTPUT_DIM; output_i++) {
+                ACCUMULATOR_TYPE d_pre_activation = d_activation_d_x<typename DEVICE::SPEC::MATH, ACCUMULATOR_TYPE, LAYER_SPEC::ACTIVATION_FUNCTION>((ACCUMULATOR_TYPE)get(layer.pre_activations, batch_i, output_i)) * (ACCUMULATOR_TYPE)get(d_output, batch_i, output_i);
+                for(TI input_j = 0; input_j < SPEC::INPUT_DIM; input_j++) {
+                    if(output_i == 0){
+                        set(d_input, batch_i, input_j, 0);
+                    }
+                    increment(d_input, batch_i, input_j, (ACCUMULATOR_TYPE)get(device, layer.weights.parameters, output_i, input_j) * d_pre_activation);
+                }
+            }
+        }
+    }
+
+    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::dense::LayerGradient<LAYER_SPEC>& layer, const Matrix<INPUT_SPEC>& input, Matrix<D_OUTPUT_SPEC>& d_output, nn::layers::dense::Buffer&, const Mode<MODE>& mode = Mode<mode::Default<>>{}) {
+        // todo: create sparate function that does not set d_input (to save cost on backward pass for the first layer)
+        // todo: think about storing gradient in column major order to avoid iterating over the minor dimension
+        static_assert(nn::layers::dense::check_input_output<LAYER_SPEC, INPUT_SPEC, D_OUTPUT_SPEC>);
+        constexpr auto INPUT_DIM = LAYER_SPEC::INPUT_DIM;
+        constexpr auto OUTPUT_DIM = LAYER_SPEC::OUTPUT_DIM;
+        constexpr auto BATCH_SIZE = D_OUTPUT_SPEC::ROWS;
+        using ACCUMULATOR_TYPE = typename LAYER_SPEC::TYPE_POLICY::template GET<numeric_types::categories::Accumulator>;
+        using TI = typename DEVICE::index_t;
+
+        for(TI batch_i=0; batch_i < BATCH_SIZE; batch_i++){
+            for(TI output_i = 0; output_i < OUTPUT_DIM; output_i++) {
+                ACCUMULATOR_TYPE d_pre_activation = d_activation_d_x<typename DEVICE::SPEC::MATH, ACCUMULATOR_TYPE, LAYER_SPEC::ACTIVATION_FUNCTION>((ACCUMULATOR_TYPE)get(layer.pre_activations, batch_i, output_i)) * (ACCUMULATOR_TYPE)get(d_output, batch_i, output_i);
+                increment(device, layer.biases.gradient, d_pre_activation, output_i);
+                for(TI input_i = 0; input_i < INPUT_DIM; input_i++){
+                    increment(device, layer.weights.gradient, d_pre_activation * (ACCUMULATOR_TYPE)get(input, batch_i, input_i), output_i, input_i);
+                }
+            }
+        }
+    }
+
+    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::dense::LayerGradient<LAYER_SPEC>& layer, const Matrix<INPUT_SPEC>& input, Matrix<D_OUTPUT_SPEC>& d_output, Matrix<D_INPUT_SPEC>& d_input, nn::layers::dense::Buffer&, const Mode<MODE>& mode = Mode<mode::Default<>>{}) {
+        // todo: create sparate function that does not set d_input (to save cost on backward pass for the first layer)
+        // todo: think about storing gradient in column major order to avoid iterating over the minor dimension
+        static_assert(nn::layers::dense::check_input_output<LAYER_SPEC, D_INPUT_SPEC, D_OUTPUT_SPEC>);
+        static_assert(nn::layers::dense::check_input_output<LAYER_SPEC, INPUT_SPEC, D_OUTPUT_SPEC>);
+        constexpr auto INPUT_DIM = LAYER_SPEC::INPUT_DIM;
+        constexpr auto OUTPUT_DIM = LAYER_SPEC::OUTPUT_DIM;
+        constexpr auto BATCH_SIZE = D_OUTPUT_SPEC::ROWS;
+        using ACCUMULATOR_TYPE = typename LAYER_SPEC::TYPE_POLICY::template GET<numeric_types::categories::Accumulator>;
+        using TI = typename DEVICE::index_t;
+
+        for(TI batch_i=0; batch_i < BATCH_SIZE; batch_i++){
+            for(TI output_i = 0; output_i < OUTPUT_DIM; output_i++) {
+                ACCUMULATOR_TYPE d_pre_activation = d_activation_d_x<typename DEVICE::SPEC::MATH, ACCUMULATOR_TYPE, LAYER_SPEC::ACTIVATION_FUNCTION>((ACCUMULATOR_TYPE)get(layer.pre_activations, batch_i, output_i)) * (ACCUMULATOR_TYPE)get(d_output, batch_i, output_i);
+                increment(device, layer.biases.gradient, d_pre_activation, output_i);
+                for(TI input_i = 0; input_i < INPUT_DIM; input_i++){
+                    if(output_i == 0){
+                        set(d_input, batch_i, input_i, 0);
+                    }
+                    increment(d_input, batch_i, input_i, (ACCUMULATOR_TYPE)get(device, layer.weights.parameters, output_i, input_i) * d_pre_activation);
+                    increment(device, layer.weights.gradient, d_pre_activation * (ACCUMULATOR_TYPE)get(input, batch_i, input_i), output_i, input_i);
+                }
+            }
+        }
+    }
+
+#endif
+    template<typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void zero_gradient(DEVICE& device, nn::layers::dense::LayerGradient<SPEC>& layer) {
+        zero_gradient(device, layer.weights);
+        zero_gradient(device, layer.biases);
+    }
+    template<typename DEVICE, typename SOURCE_SPEC, typename TARGET_SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void add_gradient(DEVICE& device, nn::layers::dense::LayerGradient<SOURCE_SPEC>& source, nn::layers::dense::LayerGradient<TARGET_SPEC>& target) {
+        add_gradient(device, source.weights, target.weights);
+        add_gradient(device, source.biases, target.biases);
+    }
+    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::dense::LayerGradient<SOURCE_SPEC>& source, nn::layers::dense::LayerGradient<TARGET_SPEC>& target) {
+        copy_gradient(source_device, target_device, source.weights, target.weights);
+        copy_gradient(source_device, target_device, source.biases, target.biases);
+    }
+    template<typename DEVICE, typename SPEC, typename OPTIMIZER>
+    RL_TOOLS_FUNCTION_PLACEMENT void update(DEVICE& device, nn::layers::dense::LayerGradient<SPEC>& layer, OPTIMIZER& optimizer){
+        update(device, layer.weights, optimizer);
+        update(device, layer.biases, optimizer);
+    }
+
+    template<typename DEVICE, typename SPEC, typename OPTIMIZER>
+    RL_TOOLS_FUNCTION_PLACEMENT void _reset_optimizer_state(DEVICE& device, nn::layers::dense::LayerGradient<SPEC>& layer, OPTIMIZER& optimizer) {
+        _reset_optimizer_state(device, layer.weights, optimizer);
+        _reset_optimizer_state(device, layer.biases, 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::dense::LayerForward<SOURCE_SPEC>& source, nn::layers::dense::LayerForward<TARGET_SPEC>& target){
+        static_assert(nn::layers::dense::check_spec_memory<SOURCE_SPEC, TARGET_SPEC>);
+        copy(source_device, target_device, source.weights, target.weights);
+        copy(source_device, target_device, source.biases, target.biases);
+    }
+    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::dense::LayerBackward<SOURCE_SPEC>& source, nn::layers::dense::LayerBackward<TARGET_SPEC>& target){
+        static_assert(nn::layers::dense::check_spec_memory<SOURCE_SPEC, TARGET_SPEC>);
+        copy(source_device, target_device, static_cast<const nn::layers::dense::LayerForward<SOURCE_SPEC>&>(source), static_cast<nn::layers::dense::LayerForward<TARGET_SPEC>&>(target));
+        copy(source_device, target_device, source.pre_activations, target.pre_activations);
+    }
+    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::dense::LayerGradient<SOURCE_SPEC>& source, nn::layers::dense::LayerGradient<TARGET_SPEC>& target){
+        static_assert(nn::layers::dense::check_spec_memory<SOURCE_SPEC, TARGET_SPEC>);
+        copy(source_device, target_device, static_cast<const nn::layers::dense::LayerBackward<SOURCE_SPEC>&>(source), static_cast<nn::layers::dense::LayerBackward<TARGET_SPEC>&>(target));
+        copy(source_device, target_device, source.output, target.output);
+
+    }
+    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::layers::dense::LayerForward<TARGET_SPEC>& target){
+        static_assert(nn::layers::dense::check_spec_memory<typename SOURCE::SPEC, TARGET_SPEC>);
+        copy_from_generic(source_device, target_device, source.weights, target.weights);
+        copy_from_generic(source_device, target_device, source.biases, target.biases);
+    }
+    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::layers::dense::LayerBackward<TARGET_SPEC>& target){
+        static_assert(nn::layers::dense::check_spec_memory<typename SOURCE::SPEC, TARGET_SPEC>);
+        copy_from_generic(source_device, target_device, static_cast<const typename SOURCE::PARENT&>(source), static_cast<nn::layers::dense::LayerForward<TARGET_SPEC>&>(target));
+        copy_from_generic(source_device, target_device, source.pre_activations, target.pre_activations);
+    }
+    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::layers::dense::LayerGradient<TARGET_SPEC>& target){
+        static_assert(nn::layers::dense::check_spec_memory<typename SOURCE::SPEC, TARGET_SPEC>);
+        copy_from_generic(source_device, target_device, static_cast<const typename SOURCE::PARENT&>(source), static_cast<nn::layers::dense::LayerBackward<TARGET_SPEC>&>(target));
+        copy_from_generic(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::dense::LayerForward<SPEC_1>* l1, const rl_tools::nn::layers::dense::LayerForward<SPEC_2>* l2) {
+        static_assert(nn::layers::dense::check_spec_memory<SPEC_1, SPEC_2>);
+        using T = typename SPEC_1::TYPE_POLICY::DEFAULT;
+        T acc = 0;
+        acc += abs_diff(device, l1->weights, l2->weights);
+        acc += abs_diff(device, l1->biases, l2->biases);
+        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::dense::LayerForward<SPEC_1>& l1, const rl_tools::nn::layers::dense::LayerForward<SPEC_2>& l2) {
+        static_assert(nn::layers::dense::check_spec_memory<SPEC_1, SPEC_2>);
+        return abs_diff(device, &l1, &l2);
+    }
+    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::dense::LayerBackward<SPEC_1>* l1, const rl_tools::nn::layers::dense::LayerBackward<SPEC_2>* l2) {
+        static_assert(nn::layers::dense::check_spec_memory<SPEC_1, SPEC_2>);
+        using T = typename SPEC_1::TYPE_POLICY::DEFAULT ;
+        T acc = abs_diff(device, (rl_tools::nn::layers::dense::LayerForward<SPEC_1>*) l1, (rl_tools::nn::layers::dense::LayerForward<SPEC_2>*) l2);
+        acc += abs_diff(device, l1->pre_activations, l2->pre_activations);
+        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::dense::LayerBackward<SPEC_1>& l1, const rl_tools::nn::layers::dense::LayerBackward<SPEC_2>& l2) {
+        static_assert(nn::layers::dense::check_spec_memory<SPEC_1, SPEC_2>);
+        return abs_diff(device, &l1, &l2);
+    }
+    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::dense::LayerGradient<SPEC_1>* l1, const rl_tools::nn::layers::dense::LayerGradient<SPEC_2>* l2) {
+        static_assert(nn::layers::dense::check_spec_memory<SPEC_1, SPEC_2>);
+        using T = typename SPEC_1::TYPE_POLICY::DEFAULT;
+        T acc = abs_diff(device, (rl_tools::nn::layers::dense::LayerBackward<SPEC_1>*) l1, (rl_tools::nn::layers::dense::LayerBackward<SPEC_2>*) l2);
+        acc += abs_diff(device, l1->output, l2->output);
+        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::dense::LayerGradient<SPEC_1>& l1, const rl_tools::nn::layers::dense::LayerGradient<SPEC_2>& l2) {
+        static_assert(nn::layers::dense::check_spec_memory<SPEC_1, SPEC_2>);
+        return abs_diff(device, &l1, &l2);
+    }
+    template <typename DEVICE>
+    RL_TOOLS_FUNCTION_PLACEMENT auto abs_diff(DEVICE& device, const rl_tools::nn::layers::dense::State& s1, const rl_tools::nn::layers::dense::State& s2) {
+        return 0;
+    }
+    template <typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void reset_forward_state(DEVICE& device, rl_tools::nn::layers::dense::LayerBackward<SPEC>& l) {
+        set_all(device, l.pre_activations, 0);
+    }
+    template <typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void reset_forward_state(DEVICE& device, rl_tools::nn::layers::dense::LayerGradient<SPEC>& l) {
+        reset_forward_state(device, static_cast<rl_tools::nn::layers::dense::LayerBackward<SPEC>&>(l));
+        set_all(device, l.output, 0);
+    }
+    template <typename DEVICE, typename SPEC, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT bool is_nan(DEVICE& device, const rl_tools::nn::layers::dense::LayerForward<SPEC>& l, const Mode<MODE>& mode = Mode<mode::Default<>>{}){
+        return is_nan(device, l.weights, mode) || is_nan(device, l.biases, mode);
+    }
+    template <typename DEVICE, typename SPEC, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT bool is_nan(DEVICE& device, const rl_tools::nn::layers::dense::LayerBackward<SPEC>& l, const Mode<MODE>& mode = Mode<mode::Default<>>{}) {
+        bool upstream_nan = is_nan(device, static_cast<const rl_tools::nn::layers::dense::LayerForward<SPEC>&>(l), mode);
+        if(mode::is<MODE, nn::parameters::mode::ParametersOnly>){
+            return upstream_nan;
+        }
+        return  upstream_nan || is_nan(device, l.pre_activations, mode);
+    }
+    template <typename DEVICE, typename SPEC, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT bool is_nan(DEVICE& device, const rl_tools::nn::layers::dense::LayerGradient<SPEC>& l, const Mode<MODE>& mode = Mode<mode::Default<>>{}){
+        bool upstream_nan = is_nan(device, static_cast<const rl_tools::nn::layers::dense::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::dense::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::dense::LayerGradient<SPEC>& l){
+        auto tensor_flat = to_tensor(device, l.output);
+        auto tensor = view_memory<typename SPEC::OUTPUT_SHAPE>(device, tensor_flat);
+        return tensor;
+    }
+}
+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::dense::LayerForward<LAYER_SPEC>& layer, const Tensor<INPUT_SPEC>& input, Tensor<OUTPUT_SPEC>& output, nn::layers::dense::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::dense::LayerForward<LAYER_SPEC>& layer, const Tensor<INPUT_SPEC>& input, nn::layers::dense::State& state, Tensor<OUTPUT_SPEC>& output, nn::layers::dense::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::dense::LayerBackward<LAYER_SPEC>& layer, const Tensor<INPUT_SPEC>& input, Tensor<OUTPUT_SPEC>& output, nn::layers::dense::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::dense::LayerGradient<LAYER_SPEC>& layer, const Tensor<INPUT_SPEC>& input, nn::layers::dense::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::dense::LayerGradient<LAYER_SPEC>& layer, const Tensor<INPUT_SPEC>& input, Tensor<OUTPUT_SPEC>& output, nn::layers::dense::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::dense::LayerBackward<LAYER_SPEC>& layer, const Tensor<D_OUTPUT_SPEC>& d_output, Tensor<D_INPUT_SPEC>& d_input, nn::layers::dense::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::dense::LayerGradient<LAYER_SPEC>& layer, const Tensor<INPUT_SPEC>& input, Tensor<D_OUTPUT_SPEC>& d_output, nn::layers::dense::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::dense::LayerGradient<LAYER_SPEC>& layer, const Tensor<INPUT_SPEC>& input, Tensor<D_OUTPUT_SPEC>& d_output, Tensor<D_INPUT_SPEC>& d_input, nn::layers::dense::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::dense::LayerGradient<SPEC>& layer) {
+        return gradient_norm(device, layer.weights) + gradient_norm(device, layer.biases);
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+
+#endif