hyperray 0.0.1__py3-none-any.whl

hyperdrone/_vendor/rl-tools/include/rl_tools/nn/layers/gru/helper_operations_generic.h

@@ -0,0 +1,251 @@
+#include "../../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_NN_LAYERS_GRU_HELPER_OPERATIONS_GENERIC_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_NN_LAYERS_GRU_HELPER_OPERATIONS_GENERIC_H
+
+#include "layer.h"
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools::nn::layers::gru::mode{
+
+        template <typename SPEC, typename MODE, typename MODE_SPEC>
+        RL_TOOLS_FUNCTION_PLACEMENT bool reset_full_batch(const Mode<nn::layers::gru::ResetMode<MODE, MODE_SPEC>>& mode){
+            return false;
+        }
+        template <typename SPEC, typename MODE>
+        RL_TOOLS_FUNCTION_PLACEMENT bool reset_full_batch(const Mode<MODE>& mode){
+            return false;
+        }
+
+        template <typename SPEC, typename MODE, typename MODE_SPEC>
+        RL_TOOLS_FUNCTION_PLACEMENT constexpr bool can_reset_sample(const Mode<nn::layers::gru::ResetMode<MODE, MODE_SPEC>>){
+            return true;
+        }
+        template <typename SPEC, typename MODE>
+        RL_TOOLS_FUNCTION_PLACEMENT constexpr bool can_reset_sample(const Mode<MODE>){
+            return false;
+        }
+
+        template <typename SPEC, typename DEVICE, typename MODE, typename MODE_SPEC, typename TI>
+        RL_TOOLS_FUNCTION_PLACEMENT bool reset_sample(DEVICE& device, const Mode<nn::layers::gru::ResetMode<MODE, MODE_SPEC>>& mode, TI step_i, TI batch_sample_i){
+            return get(device, mode.reset_container, (typename MODE_SPEC::TI)step_i, (typename MODE_SPEC::TI)batch_sample_i, (typename MODE_SPEC::TI)0);
+        }
+        template <typename SPEC, typename DEVICE, typename MODE, typename TI>
+        RL_TOOLS_FUNCTION_PLACEMENT bool reset_sample(DEVICE& device, const Mode<MODE>& mode, TI step_i, TI batch_sample_i){
+            return false;
+        }
+
+        template <typename DEVICE, typename MODE, typename MODE_SPEC>
+        RL_TOOLS_FUNCTION_PLACEMENT typename DEVICE::index_t num_resets(DEVICE& device, const Mode<nn::layers::gru::ResetMode<MODE, MODE_SPEC>>& mode){
+            return cast_reduce_sum<typename MODE_SPEC::TI>(device, mode.reset_container);
+        }
+        template <typename DEVICE, typename MODE, bool SHOULD_NEVER=false>
+        RL_TOOLS_FUNCTION_PLACEMENT typename DEVICE::index_t num_resets(DEVICE& device, const Mode<MODE>& mode){
+            static_assert(SHOULD_NEVER, "This should not be called");
+            return 1;
+        }
+    }
+namespace rl_tools::nn::layers::gru::helper{
+    template<typename DEVICE, typename SPEC_1, typename SPEC_2, typename SPEC_BIAS, typename SPEC_OUT>
+    RL_TOOLS_FUNCTION_PLACEMENT void matrix_multiply_transpose_bias(DEVICE& device, const Tensor<SPEC_1>& t1, const Tensor<SPEC_2>& t2, const Tensor<SPEC_BIAS>& bias, Tensor<SPEC_OUT>& result){
+#ifdef RL_TOOLS_ENABLE_TRACY
+        ZoneScopedN("gru::matrix_multiply_transpose_bias");
+#endif
+        // Y = WX
+        // Y^T = X^T W^T
+        // W = t1, X^T = t2, Y^T = result
+        // Y^T = result = t2 t1^T
+        static_assert(length(typename SPEC_1::SHAPE{}) == 2);
+        static_assert(length(typename SPEC_2::SHAPE{}) == 2);
+        static_assert(length(typename SPEC_OUT::SHAPE{}) == 2);
+        static_assert(get<1>(typename SPEC_1::SHAPE{}) == get<1>(typename SPEC_2::SHAPE{})); // INPUT_DIM
+        static_assert(get<0>(typename SPEC_2::SHAPE{}) == get<0>(typename SPEC_OUT::SHAPE{})); // BATCH_SIZE
+        static_assert(get<0>(typename SPEC_1::SHAPE{}) == get<1>(typename SPEC_OUT::SHAPE{})); // HIDDEN_DIM
+        static_assert(length(typename SPEC_BIAS::SHAPE{}) == 1);
+        static_assert(get<0>(typename SPEC_BIAS::SHAPE{}) == get<0>(typename SPEC_1::SHAPE{}));
+        using T = typename SPEC_1::T;
+        using TI = typename DEVICE::index_t;
+        for(TI i=0; i < get<0>(typename SPEC_OUT::SHAPE{}); i++){
+            for(TI j=0; j < get<1>(typename SPEC_OUT::SHAPE{}); j++){
+                T bias_value = get(device, bias, j);
+                set(device, result, bias_value, i, j);
+            }
+        }
+        auto t1_transpose = permute(device, t1, tensor::PermutationSpec<1, 0>{});
+        matrix_multiply_accumulate(device, t2, t1_transpose, result);
+//        for(TI i=0; i < get<0>(typename SPEC_1::SHAPE{}); ++i){
+//            for(TI j=0; j < get<0>(typename SPEC_2::SHAPE{}); ++j){
+//                T acc = get(device, bias, i);
+//                for(TI k=0; k < get<1>(typename SPEC_1::SHAPE{}); ++k){
+//                    acc += get(device, t1, i, k) * get(device, t2, j, k);
+//                }
+//                set(device, result, acc, j, i);
+//            }
+//        }
+    }
+    template<typename DEVICE, typename SPEC_1, typename SPEC_2, typename SPEC_BIAS, typename SPEC_OUT>
+    RL_TOOLS_FUNCTION_PLACEMENT void matrix_multiply_transpose_bias_accumulate(DEVICE& device, const Tensor<SPEC_1>& t1, const Tensor<SPEC_2>& t2, const Tensor<SPEC_BIAS>& bias, Tensor<SPEC_OUT>& result){
+#ifdef RL_TOOLS_ENABLE_TRACY
+        ZoneScopedN("gru::matrix_multiply_transpose_bias_accumulate");
+#endif
+        // Y^T = WX^T
+        static_assert(length(typename SPEC_1::SHAPE{}) == 2);
+        static_assert(length(typename SPEC_2::SHAPE{}) == 2);
+        static_assert(length(typename SPEC_OUT::SHAPE{}) == 2);
+        static_assert(get<1>(typename SPEC_1::SHAPE{}) == get<1>(typename SPEC_2::SHAPE{})); // INPUT_DIM
+        static_assert(get<0>(typename SPEC_2::SHAPE{}) == get<0>(typename SPEC_OUT::SHAPE{})); // BATCH_SIZE
+        static_assert(get<0>(typename SPEC_1::SHAPE{}) == get<1>(typename SPEC_OUT::SHAPE{})); // HIDDEN_DIM
+        static_assert(length(typename SPEC_BIAS::SHAPE{}) == 1);
+        static_assert(get<0>(typename SPEC_BIAS::SHAPE{}) == get<0>(typename SPEC_1::SHAPE{}));
+        using T = typename SPEC_1::T;
+        using TI = typename DEVICE::index_t;
+        for(TI i=0; i < get<0>(typename SPEC_OUT::SHAPE{}); i++){
+            for(TI j=0; j < get<1>(typename SPEC_OUT::SHAPE{}); j++){
+                T value = get(device, result, i, j) + get(device, bias, j);
+                set(device, result, value, i, j);
+            }
+        }
+        auto t1_transpose = permute(device, t1, tensor::PermutationSpec<1, 0>{});
+        matrix_multiply_accumulate(device, t2, t1_transpose, result);
+    }
+
+    template<typename DEVICE, typename SPEC_1, typename SPEC_2, typename SPEC_BIAS, typename SPEC_OUT>
+    RL_TOOLS_FUNCTION_PLACEMENT void matrix_multiply_broadcast_transpose_bias(DEVICE& device, const Tensor<SPEC_1>& t1, const Tensor<SPEC_2>& t2, const Tensor<SPEC_BIAS>& bias, Tensor<SPEC_OUT>& result){
+        // t2_expand = expand_along_cols(t2)
+        // result^T = t1 t2_expand + bias
+#ifdef RL_TOOLS_ENABLE_TRACY
+        ZoneScopedN("gru::matrix_multiply_broadcast_transpose_bias");
+#endif
+        // Y^T = WX^T
+        static_assert(length(typename SPEC_1::SHAPE{}) == 2);
+        static_assert(length(typename SPEC_2::SHAPE{}) == 1);
+        static_assert(length(typename SPEC_OUT::SHAPE{}) == 2);
+        static_assert(get<1>(typename SPEC_1::SHAPE{}) == get<0>(typename SPEC_2::SHAPE{})); // INPUT_DIM
+//        static_assert(get<0>(typename SPEC_2::SHAPE{}) == get<0>(typename SPEC_OUT::SHAPE{})); // BATCH_SIZE
+        static_assert(get<0>(typename SPEC_1::SHAPE{}) == get<1>(typename SPEC_OUT::SHAPE{})); // HIDDEN_DIM
+        static_assert(length(typename SPEC_BIAS::SHAPE{}) == 1);
+        static_assert(get<0>(typename SPEC_BIAS::SHAPE{}) == get<0>(typename SPEC_1::SHAPE{}));
+        using T = typename SPEC_1::T;
+        using TI = typename DEVICE::index_t;
+        for(TI i=0; i < get<1>(typename SPEC_OUT::SHAPE{}); ++i){
+            for(TI j=0; j < get<0>(typename SPEC_OUT::SHAPE{}); ++j){
+                T acc = get(device, bias, i);
+                for(TI k=0; k < get<1>(typename SPEC_1::SHAPE{}); ++k){
+                    acc += get(device, t1, i, k) * get(device, t2, k);
+                }
+                set(device, result, acc, j, i);
+            }
+        }
+    }
+
+    namespace unary_operations{
+        template <typename DEVICE, typename PARAMETER, typename T>
+        RL_TOOLS_FUNCTION_PLACEMENT T d_sigmoid(DEVICE& device, const PARAMETER& parameter, T a){
+            T s = sigmoid(device, parameter, a);
+            return s * (1 - s);
+        }
+        template <typename DEVICE, typename PARAMETER, typename T>
+        RL_TOOLS_FUNCTION_PLACEMENT T d_sigmoid_post_activation(DEVICE& device, const PARAMETER& parameter, T s){
+            return s * (1 - s);
+        }
+        template <typename DEVICE, typename PARAMETER, typename T>
+        RL_TOOLS_FUNCTION_PLACEMENT T d_tanh(DEVICE& device, const PARAMETER& parameter, T a){
+            T t = tanh(device, parameter, a);// todo: make version that can take advantage of a stored value so that tanh does not have to be calculated again
+            return 1 - t * t;
+        }
+        template <typename DEVICE, typename PARAMETER, typename T>
+        RL_TOOLS_FUNCTION_PLACEMENT T d_tanh_post_activation(DEVICE& device, const PARAMETER& parameter, T t){
+            return 1 - t * t;
+        }
+    }
+    template<typename DEVICE, typename SPEC, typename SPEC_OUTPUT>
+    RL_TOOLS_FUNCTION_PLACEMENT void d_sigmoid(DEVICE& device, Tensor<SPEC>& t, Tensor<SPEC_OUTPUT>& output){
+        static_assert(tensor::same_dimensions<SPEC, SPEC_OUTPUT>());
+        using T = typename SPEC::T;
+        unary_operation(device, tensor::OperationLegacy<unary_operations::d_sigmoid<DEVICE, tensor::OperationEmptyParameter, T>, tensor::OperationEmptyParameter>{}, t, output);
+    }
+    template<typename DEVICE, typename SPEC, typename SPEC_OUTPUT>
+    RL_TOOLS_FUNCTION_PLACEMENT void d_sigmoid_post_activation(DEVICE& device, Tensor<SPEC>& t){
+        // this function takes sigmoid(x) as an input instead of x (to reuse it in the calculation)
+        static_assert(tensor::same_dimensions<SPEC, SPEC_OUTPUT>());
+        using T = typename SPEC::T;
+        unary_operation(device, tensor::OperationLegacy<unary_operations::d_sigmoid_post_activation<DEVICE, tensor::OperationEmptyParameter, T>, tensor::OperationEmptyParameter>{}, t);
+    }
+    template<typename DEVICE, typename SPEC, typename SPEC_OUTPUT>
+    RL_TOOLS_FUNCTION_PLACEMENT void d_sigmoid_post_activation(DEVICE& device, Tensor<SPEC>& t, Tensor<SPEC_OUTPUT>& output){
+        // this function takes sigmoid(x) as an input instead of x (to reuse it in the calculation)
+        static_assert(tensor::same_dimensions<SPEC, SPEC_OUTPUT>());
+        using T = typename SPEC::T;
+        unary_operation(device, tensor::OperationLegacy<unary_operations::d_sigmoid_post_activation<DEVICE, tensor::OperationEmptyParameter, T>, tensor::OperationEmptyParameter>{}, t, output);
+    }
+    template<typename DEVICE, typename SPEC, typename SPEC_OUTPUT>
+    RL_TOOLS_FUNCTION_PLACEMENT void d_tanh(DEVICE& device, Tensor<SPEC>& t, Tensor<SPEC_OUTPUT>& output){
+        static_assert(tensor::same_dimensions<SPEC, SPEC_OUTPUT>());
+        using T = typename SPEC::T;
+        unary_operation(device, tensor::OperationLegacy<unary_operations::d_tanh<DEVICE, tensor::OperationEmptyParameter, T>, tensor::OperationEmptyParameter>{}, t, output);
+    }
+    template<typename DEVICE, typename SPEC, typename SPEC_OUTPUT>
+    RL_TOOLS_FUNCTION_PLACEMENT void d_tanh_post_activation(DEVICE& device, Tensor<SPEC>& t, Tensor<SPEC_OUTPUT>& output){
+        // this function takes tanh(x) as an input instead of x (to reuse it in the calculation)
+        static_assert(tensor::same_dimensions<SPEC, SPEC_OUTPUT>());
+        using T = typename SPEC::T;
+        unary_operation(device, tensor::OperationLegacy<unary_operations::d_tanh_post_activation<DEVICE, tensor::OperationEmptyParameter, T>, tensor::OperationEmptyParameter>{}, t, output);
+    }
+
+    template<typename LAYER_SPEC, typename DEVICE, typename INITIAL_HIDDEN_STATE_SPEC, typename OUTPUT_PREVIOUS_STEP_SPEC, typename PREVIOUS_OUTPUT_SCRATCH_SPEC, typename MODE>
+    RL_TOOLS_FUNCTION_PLACEMENT void select_previous_output(DEVICE& device, const Tensor<INITIAL_HIDDEN_STATE_SPEC>& initial_hidden_state, const Tensor<OUTPUT_PREVIOUS_STEP_SPEC>& output_previous_step, Tensor<PREVIOUS_OUTPUT_SCRATCH_SPEC>& previous_output_scratch, typename DEVICE::index_t step_i, const Mode<MODE>& mode){
+        using TI = typename DEVICE::index_t;
+        constexpr TI BATCH_SIZE = get<0>(typename PREVIOUS_OUTPUT_SCRATCH_SPEC::SHAPE{});
+        for(TI sample_i = 0; sample_i < BATCH_SIZE; sample_i++){
+            auto target = view(device, previous_output_scratch, sample_i);
+            if(nn::layers::gru::mode::reset_sample<LAYER_SPEC>(device, mode, step_i, sample_i)){
+                copy(device, device, initial_hidden_state, target);
+            }
+            else{
+                auto source = view(device, output_previous_step, sample_i);
+                copy(device, device, source, target);
+            }
+        }
+    }
+    template<typename LAYER_SPEC, bool ACCUMULATE_INITIAL_HIDDEN_STATE_GRADIENT, typename DEVICE, typename D_OUTPUT_STEP_SPEC, typename Z_POST_ACTIVATION_SPEC, typename D_OUTPUT_PREVIOUS_STEP_SPEC, typename INITIAL_HIDDEN_STATE_GRADIENT_SPEC, typename MODE>
+    RL_TOOLS_FUNCTION_PLACEMENT void multiply_accumulate_d_output_previous_step(DEVICE& device, const Tensor<D_OUTPUT_STEP_SPEC>& d_output_step, const Tensor<Z_POST_ACTIVATION_SPEC>& z_post_activation, Tensor<D_OUTPUT_PREVIOUS_STEP_SPEC>& d_output_previous_step, Tensor<INITIAL_HIDDEN_STATE_GRADIENT_SPEC>& initial_hidden_state_gradient, typename DEVICE::index_t step_i, const Mode<MODE>& mode){
+        using TI = typename DEVICE::index_t;
+        constexpr TI BATCH_SIZE = get<0>(typename D_OUTPUT_STEP_SPEC::SHAPE{});
+        for(TI sample_i = 0; sample_i < BATCH_SIZE; sample_i++){
+            auto d_output_step_sample = view(device, d_output_step, sample_i);
+            auto z_post_activation_sample = view(device, z_post_activation, sample_i);
+            if(nn::layers::gru::mode::reset_sample<LAYER_SPEC>(device, mode, step_i, sample_i)){
+                if constexpr(ACCUMULATE_INITIAL_HIDDEN_STATE_GRADIENT){
+                    multiply_accumulate(device, d_output_step_sample, z_post_activation_sample, initial_hidden_state_gradient);
+                }
+            }
+            else{
+                auto d_output_previous_step_sample = view(device, d_output_previous_step, sample_i);
+                multiply_accumulate(device, d_output_step_sample, z_post_activation_sample, d_output_previous_step_sample);
+            }
+        }
+    }
+    namespace detail{
+        template<typename LAYER_SPEC, bool RESET_SAMPLES, typename DEVICE, typename SPEC, typename MODE>
+        RL_TOOLS_FUNCTION_PLACEMENT void zero_samples(DEVICE& device, Tensor<SPEC>& tensor, typename DEVICE::index_t step_i, const Mode<MODE>& mode){
+            using TI = typename DEVICE::index_t;
+            constexpr TI BATCH_SIZE = get<0>(typename SPEC::SHAPE{});
+            for(TI sample_i = 0; sample_i < BATCH_SIZE; sample_i++){
+                if(nn::layers::gru::mode::reset_sample<LAYER_SPEC>(device, mode, step_i, sample_i) == RESET_SAMPLES){
+                    auto sample = view(device, tensor, sample_i);
+                    set_all(device, sample, 0);
+                }
+            }
+        }
+    }
+    template<typename LAYER_SPEC, typename DEVICE, typename SPEC, typename MODE>
+    RL_TOOLS_FUNCTION_PLACEMENT void zero_reset_samples(DEVICE& device, Tensor<SPEC>& tensor, typename DEVICE::index_t step_i, const Mode<MODE>& mode){
+        detail::zero_samples<LAYER_SPEC, true>(device, tensor, step_i, mode);
+    }
+    template<typename LAYER_SPEC, typename DEVICE, typename SPEC, typename MODE>
+    RL_TOOLS_FUNCTION_PLACEMENT void zero_continuing_samples(DEVICE& device, Tensor<SPEC>& tensor, typename DEVICE::index_t step_i, const Mode<MODE>& mode){
+        detail::zero_samples<LAYER_SPEC, false>(device, tensor, step_i, mode);
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+#endif

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

@@ -0,0 +1,206 @@
+#include "../../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_NN_LAYERS_GRU_LAYER_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_NN_LAYERS_GRU_LAYER_H
+#include "../../../mode/mode.h"
+#include "../../../nn/activation_functions.h"
+#include "../../../nn/parameters/parameters.h"
+#include "../../../nn/capability/capability.h"
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools::nn::layers::gru{
+    template<typename T_TYPE_POLICY, typename T_TI, T_TI T_HIDDEN_DIM, typename T_PARAMETER_GROUP=parameters::groups::Normal, bool T_FAST_TANH = false>
+    struct Configuration{
+        using TYPE_POLICY = T_TYPE_POLICY;
+        using TI = T_TI;
+        static constexpr T_TI HIDDEN_DIM = T_HIDDEN_DIM;
+        using PARAMETER_GROUP = T_PARAMETER_GROUP;
+        static constexpr bool FAST_TANH = T_FAST_TANH;
+        static constexpr bool LEARN_INITIAL_HIDDEN_STATE = false;
+        // Summary
+//        static constexpr auto NUM_WEIGHTS = HIDDEN_DIM * INPUT_DIM + HIDDEN_DIM; // todo
+    };
+    template <typename T_CONFIG, typename T_CAPABILITY, typename T_INPUT_SHAPE>
+    struct Specification: T_CAPABILITY, T_CONFIG{
+        using CONFIG = T_CONFIG;
+        using CAPABILITY = T_CAPABILITY;
+        using INPUT_SHAPE = T_INPUT_SHAPE;
+        using TYPE_POLICY = typename CONFIG::TYPE_POLICY;
+        using TI = typename CONFIG::TI;
+        static_assert(length(INPUT_SHAPE{}) == 3, "The input shape of the GRU must be 3 dimensional for now (sequence x batch x features)");
+        static constexpr TI INPUT_DIM = get_last(INPUT_SHAPE{});
+        static constexpr TI SEQUENCE_LENGTH = get<length(INPUT_SHAPE{})-3>(INPUT_SHAPE{});
+        template <typename NEW_INPUT_SHAPE>
+        using OUTPUT_SHAPE_FACTORY = tensor::Replace<NEW_INPUT_SHAPE, CONFIG::HIDDEN_DIM, length(NEW_INPUT_SHAPE{})-1>;
+        using OUTPUT_SHAPE = OUTPUT_SHAPE_FACTORY<INPUT_SHAPE>;
+        static constexpr TI INTERNAL_BATCH_SIZE = get<1>(INPUT_SHAPE{}); // Since the Dense layer is based on Matrices (2D Tensors) the dense layer operation is broadcasted over the leading dimensions. Hence, the actual batch size is the product of all leading dimensions, excluding the last one (containing the features). Since rl_tools::matrix_view is used for zero-cost conversion the INTERNAL_BATCH_SIZE accounts for all leading dimensions.
+        static constexpr TI NUM_WEIGHTS = CONFIG::HIDDEN_DIM * (INPUT_DIM + 1) * 3 + CONFIG::HIDDEN_DIM * (CONFIG::HIDDEN_DIM+1) * 3;
+    };
+
+    namespace buffers{
+        template <typename T_SPEC, bool T_DYNAMIC_ALLOCATION>
+        struct Specification{
+            using SPEC = T_SPEC;
+            static constexpr bool DYNAMIC_ALLOCATION = T_DYNAMIC_ALLOCATION;
+        };
+        template <typename T_SPEC>
+        struct Evaluation{
+            using BUFFER_SPEC = T_SPEC;
+            using GRU_SPEC = typename BUFFER_SPEC::SPEC;
+            using TYPE_POLICY = typename GRU_SPEC::TYPE_POLICY;
+            using TI = typename GRU_SPEC::TI;
+            static constexpr TI BATCH_SIZE = GRU_SPEC::INTERNAL_BATCH_SIZE;
+            using T_ACC = typename TYPE_POLICY::template GET<numeric_types::categories::Accumulator>;
+            using POST_ACTIVATION_SPEC = tensor::Specification<T_ACC, TI, tensor::Shape<TI, BATCH_SIZE, 3*GRU_SPEC::HIDDEN_DIM>, BUFFER_SPEC::DYNAMIC_ALLOCATION>;
+            Tensor<POST_ACTIVATION_SPEC> post_activation;
+            using N_PRE_PRE_ACTIVATION_SPEC = tensor::Specification<T_ACC, TI, tensor::Shape<TI, BATCH_SIZE, GRU_SPEC::HIDDEN_DIM>, BUFFER_SPEC::DYNAMIC_ALLOCATION>;
+            Tensor<N_PRE_PRE_ACTIVATION_SPEC> n_pre_pre_activation;
+            using STEP_BY_STEP_OUTPUT_SPEC = tensor::Specification<T_ACC, TI, tensor::Shape<TI, 1, BATCH_SIZE, GRU_SPEC::HIDDEN_DIM>, BUFFER_SPEC::DYNAMIC_ALLOCATION>;
+            Tensor<STEP_BY_STEP_OUTPUT_SPEC> step_by_step_output;
+
+            using PREVIOUS_OUTPUT_SCRATCH_SPEC = tensor::Specification<T_ACC, TI, tensor::Shape<TI, BATCH_SIZE, GRU_SPEC::HIDDEN_DIM>, BUFFER_SPEC::DYNAMIC_ALLOCATION>;
+            Tensor<PREVIOUS_OUTPUT_SCRATCH_SPEC> previous_output_scratch;
+        };
+    }
+
+    template <typename TI, typename RESET_CONTAINER>
+    using ResetModeSpecification = mode::sequential::ResetModeSpecification<TI, RESET_CONTAINER>;
+    template <typename BASE, typename SPEC>
+    using ResetMode = mode::sequential::ResetMode<BASE, SPEC>;
+    template <typename BASE, typename SPEC = bool>
+    using NoAutoResetMode = mode::sequential::NoAutoResetMode<BASE, SPEC>;
+
+    template <typename T_SPEC, bool T_DYNAMIC_ALLOCATION>
+    struct StateSpecification{
+        using SPEC = T_SPEC;
+        static constexpr bool DYNAMIC_ALLOCATION = T_DYNAMIC_ALLOCATION;
+    };
+    template <typename T_SPEC>
+    struct State{
+        using STATE_SPEC = T_SPEC;
+        using SPEC = typename STATE_SPEC::SPEC;
+        using TYPE_POLICY = typename SPEC::TYPE_POLICY;
+        using TI = typename SPEC::TI;
+
+        using T_ACC = typename TYPE_POLICY::template GET<numeric_types::categories::Accumulator>;
+
+        using STATE_SHAPE = tensor::Shape<TI, SPEC::INTERNAL_BATCH_SIZE, SPEC::HIDDEN_DIM>;
+        using STATE_CONTAINER_SPEC = tensor::Specification<T_ACC, TI, STATE_SHAPE, STATE_SPEC::DYNAMIC_ALLOCATION>;
+        using STATE_CONTAINER_TYPE = Tensor<STATE_CONTAINER_SPEC>;
+        STATE_CONTAINER_TYPE state;
+        using STEP_SHAPE = tensor::Shape<TI, SPEC::INTERNAL_BATCH_SIZE>;
+        using STEP_CONTAINER_SPEC = tensor::Specification<TI, TI, STEP_SHAPE, STATE_SPEC::DYNAMIC_ALLOCATION>;
+        using STEP_CONTAINER_TYPE = Tensor<STEP_CONTAINER_SPEC>;
+        STEP_CONTAINER_TYPE step;
+    };
+
+    template<typename T_SPEC>
+    struct LayerForward{
+        using SPEC = T_SPEC;
+        using TYPE_POLICY = typename SPEC::TYPE_POLICY;
+        using TI = typename SPEC::TI;
+        static constexpr TI SEQUENCE_LENGTH = SPEC::SEQUENCE_LENGTH;
+        static constexpr TI INPUT_DIM = SPEC::INPUT_DIM;
+        static constexpr TI HIDDEN_DIM = SPEC::HIDDEN_DIM;
+        static constexpr TI OUTPUT_DIM = SPEC::HIDDEN_DIM;
+        using INPUT_SHAPE = typename SPEC::INPUT_SHAPE;
+        using OUTPUT_SHAPE = typename SPEC::OUTPUT_SHAPE;
+        template <typename NEW_INPUT_SHAPE>
+        using OUTPUT_SHAPE_FACTORY = typename SPEC::template OUTPUT_SHAPE_FACTORY<NEW_INPUT_SHAPE>;
+        using WEIGHTS_INPUT_SHAPE = tensor::Shape<TI, 3*HIDDEN_DIM, INPUT_DIM>;
+        using WEIGHTS_INPUT_PARAMETER_SPEC = typename SPEC::PARAMETER_TYPE::template Specification<TYPE_POLICY, TI, WEIGHTS_INPUT_SHAPE, typename SPEC::PARAMETER_GROUP, nn::parameters::categories::Weights, SPEC::DYNAMIC_ALLOCATION, SPEC::CONST>;
+        typename SPEC::PARAMETER_TYPE::template Instance<WEIGHTS_INPUT_PARAMETER_SPEC> weights_input;
+
+        using BIASES_INPUT_SHAPE = tensor::Shape<TI, 3*HIDDEN_DIM>;
+        using BIASES_INPUT_PARAMETER_SPEC = typename SPEC::PARAMETER_TYPE::template Specification<TYPE_POLICY, TI, BIASES_INPUT_SHAPE, typename SPEC::PARAMETER_GROUP, nn::parameters::categories::Biases, SPEC::DYNAMIC_ALLOCATION, SPEC::CONST>;
+        typename SPEC::PARAMETER_TYPE::template Instance<BIASES_INPUT_PARAMETER_SPEC> biases_input;
+
+        using WEIGHTS_HIDDEN_SHAPE = tensor::Shape<TI, 3*HIDDEN_DIM, HIDDEN_DIM>;
+        using WEIGHTS_HIDDEN_PARAMETER_SPEC = typename SPEC::PARAMETER_TYPE::template Specification<TYPE_POLICY, TI, WEIGHTS_HIDDEN_SHAPE, typename SPEC::PARAMETER_GROUP, nn::parameters::categories::Weights, SPEC::DYNAMIC_ALLOCATION, SPEC::CONST>;
+        typename SPEC::PARAMETER_TYPE::template Instance<WEIGHTS_HIDDEN_PARAMETER_SPEC> weights_hidden;
+
+        using BIASES_HIDDEN_SHAPE = tensor::Shape<TI, 3*HIDDEN_DIM>;
+        using BIASES_HIDDEN_PARAMETER_SPEC = typename SPEC::PARAMETER_TYPE::template Specification<TYPE_POLICY, TI, BIASES_HIDDEN_SHAPE, typename SPEC::PARAMETER_GROUP, nn::parameters::categories::Biases, SPEC::DYNAMIC_ALLOCATION, SPEC::CONST>;
+        typename SPEC::PARAMETER_TYPE::template Instance<BIASES_HIDDEN_PARAMETER_SPEC> biases_hidden;
+
+        using INITIAL_HIDDEN_STATE_SHAPE = tensor::Shape<TI, HIDDEN_DIM>;
+        using INITIAL_HIDDEN_STATE_PARAMETER_SPEC = typename SPEC::PARAMETER_TYPE::template Specification<TYPE_POLICY, TI, INITIAL_HIDDEN_STATE_SHAPE, typename SPEC::PARAMETER_GROUP, nn::parameters::categories::Biases, SPEC::DYNAMIC_ALLOCATION, SPEC::CONST>;
+        typename SPEC::PARAMETER_TYPE::template Instance<INITIAL_HIDDEN_STATE_PARAMETER_SPEC> initial_hidden_state;
+
+        template<bool DYNAMIC_ALLOCATION=true>
+        using State = State<StateSpecification<SPEC, DYNAMIC_ALLOCATION>>;
+        template<bool DYNAMIC_ALLOCATION=true>
+        using Buffer = buffers::Evaluation<buffers::Specification<SPEC, DYNAMIC_ALLOCATION>>;
+    };
+
+    namespace buffers{
+        template <typename T_SPEC>
+        struct Backward: Evaluation<T_SPEC>{
+            using BUFFER_SPEC = T_SPEC;
+            using LAYER_SPEC = typename T_SPEC::SPEC;
+            using TYPE_POLICY = typename LAYER_SPEC::TYPE_POLICY;
+            using TI = typename LAYER_SPEC::TI;
+            static constexpr TI INTERNAL_BATCH_SIZE = LAYER_SPEC::INTERNAL_BATCH_SIZE;
+            using T_ACC = typename TYPE_POLICY::template GET<numeric_types::categories::Accumulator>;
+            using TENSOR_SPEC = tensor::Specification<T_ACC, TI, tensor::Shape<TI, INTERNAL_BATCH_SIZE, 3*LAYER_SPEC::HIDDEN_DIM>, BUFFER_SPEC::DYNAMIC_ALLOCATION>;
+            using BUFFER_TYPE = Tensor<TENSOR_SPEC>;
+            BUFFER_TYPE buffer, buffer2;
+            typename decltype(buffer)::template VIEW_RANGE<tensor::ViewSpec<1, 2*LAYER_SPEC::HIDDEN_DIM>> buffer_rz;
+            typename decltype(buffer)::template VIEW_RANGE<tensor::ViewSpec<1, LAYER_SPEC::HIDDEN_DIM>> buffer_r, buffer_z, buffer_n;
+        };
+    }
+
+    template<typename T_SPEC>
+    struct LayerBackward: LayerForward<T_SPEC>{
+        using PARENT = LayerForward<T_SPEC>;
+        using SPEC = T_SPEC;
+        using TYPE_POLICY = typename SPEC::TYPE_POLICY;
+        using TI = typename SPEC::TI;
+        using T_ACC = typename TYPE_POLICY::template GET<numeric_types::categories::Accumulator>;
+        using FULL_HIDDEN_SHAPE = tensor::Shape<TI, SPEC::SEQUENCE_LENGTH, SPEC::INTERNAL_BATCH_SIZE, 3*SPEC::HIDDEN_DIM>;
+        using FULL_HIDDEN_SPEC = tensor::Specification<T_ACC, TI, FULL_HIDDEN_SHAPE, SPEC::DYNAMIC_ALLOCATION, tensor::RowMajorStride<FULL_HIDDEN_SHAPE>, SPEC::CONST>;
+        using FULL_HIDDEN_TYPE = Tensor<FULL_HIDDEN_SPEC>;
+        FULL_HIDDEN_TYPE post_activation;
+        using SINGLE_HIDDEN_SHAPE = tensor::Shape<TI, SPEC::SEQUENCE_LENGTH, SPEC::INTERNAL_BATCH_SIZE, SPEC::HIDDEN_DIM>;
+        using HIDDEN_SPEC = tensor::Specification<T_ACC, TI, SINGLE_HIDDEN_SHAPE, SPEC::DYNAMIC_ALLOCATION, tensor::RowMajorStride<SINGLE_HIDDEN_SHAPE>, SPEC::CONST>;
+        using HIDDEN_TYPE = Tensor<HIDDEN_SPEC>;
+        HIDDEN_TYPE n_pre_pre_activation;
+        using OUTPUT_SPEC = tensor::Specification<T_ACC, TI, SINGLE_HIDDEN_SHAPE, SPEC::DYNAMIC_ALLOCATION, tensor::RowMajorStride<SINGLE_HIDDEN_SHAPE>, SPEC::CONST>;
+        using OUTPUT_TYPE = Tensor<OUTPUT_SPEC>;
+        OUTPUT_TYPE output;
+    };
+    template<typename T_SPEC>
+    struct LayerGradient: LayerBackward<T_SPEC>{
+        using PARENT = LayerBackward<T_SPEC>;
+        using TI = typename T_SPEC::TI;
+        template<bool DYNAMIC_ALLOCATION=true>
+        using Buffer = buffers::Backward<buffers::Specification<T_SPEC, DYNAMIC_ALLOCATION>>;
+    };
+
+    template <typename SPEC, typename INPUT_SPEC, typename OUTPUT_SPEC>
+    struct assert_input_output{
+        static_assert(length(typename INPUT_SPEC::SHAPE{}) == 3 && length(typename OUTPUT_SPEC::SHAPE{}) == 3);
+        static_assert(get<0>(typename INPUT_SPEC::SHAPE{}) == get<0>(typename OUTPUT_SPEC::SHAPE{}));
+        static_assert(get<1>(typename INPUT_SPEC::SHAPE{}) == get<1>(typename OUTPUT_SPEC::SHAPE{}));
+        static_assert(get<2>(typename INPUT_SPEC::SHAPE{}) == SPEC::INPUT_DIM && get<2>(typename OUTPUT_SPEC::SHAPE{}) == SPEC::HIDDEN_DIM);
+    };
+    template <typename SPEC, typename INPUT_SPEC, typename OUTPUT_SPEC>
+    bool constexpr check_input_output = length(typename INPUT_SPEC::SHAPE{}) == 3 && length(typename OUTPUT_SPEC::SHAPE{}) == 3 &&
+            get<0>(typename INPUT_SPEC::SHAPE{}) <= get<0>(typename OUTPUT_SPEC::SHAPE{}) &&
+            get<1>(typename INPUT_SPEC::SHAPE{}) <= get<1>(typename OUTPUT_SPEC::SHAPE{}) &&
+            get<2>(typename INPUT_SPEC::SHAPE{}) == SPEC::INPUT_DIM && get<2>(typename OUTPUT_SPEC::SHAPE{}) == SPEC::HIDDEN_DIM;
+
+    template<typename CONFIG, typename CAPABILITY, typename INPUT_SHAPE>
+    using Layer =
+    typename utils::typing::conditional_t<CAPABILITY::TAG == nn::LayerCapability::Forward, LayerForward<Specification<CONFIG, CAPABILITY, INPUT_SHAPE>>,
+    typename utils::typing::conditional_t<CAPABILITY::TAG == nn::LayerCapability::Backward, LayerBackward<Specification<CONFIG, CAPABILITY, INPUT_SHAPE>>,
+    typename utils::typing::conditional_t<CAPABILITY::TAG == nn::LayerCapability::Gradient, LayerGradient<Specification<CONFIG, CAPABILITY, INPUT_SHAPE>>, void>>>;
+
+    template <typename T_CONFIG>
+    struct BindConfiguration{
+        template <typename CAPABILITY, typename INPUT_SHAPE>
+        using Layer = nn::layers::gru::Layer<T_CONFIG, CAPABILITY, INPUT_SHAPE>;
+    };
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+#endif