hyperverse 0.0.1__py3-none-any.whl

hyperdrone/_vendor/rl-tools/include/rl_tools/nn/activation_functions.h

@@ -0,0 +1,113 @@
+#include "../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_NN_ACTIVATION_FUNCTIONS_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_NN_ACTIVATION_FUNCTIONS_H
+#include "../rl_tools.h"
+#include "../devices/devices.h"
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools::nn::activation_functions {
+    enum ActivationFunction{
+        IDENTITY,
+        RELU,
+        GELU,
+        TANH,
+        FAST_TANH,
+        SIGMOID
+    };
+    template<enum ActivationFunction F>
+    constexpr bool check_activation_function = F == IDENTITY || F == RELU || F == GELU || F == TANH || F == FAST_TANH || F == SIGMOID;
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools{
+    template<typename DEVICE, typename T, nn::activation_functions::ActivationFunction F>
+    RL_TOOLS_FUNCTION_PLACEMENT inline T activation(T x){
+        using namespace nn::activation_functions;
+        static_assert(DEVICE::TYPE == devices::Type::math, "DEVICE is not a math device");
+        static_assert(nn::activation_functions::check_activation_function<F>, "Invalid activation function");
+        if constexpr(F == IDENTITY){
+            return x;
+        }
+        else if constexpr(F == RELU){
+            return math::max(DEVICE(), x, (T)0);
+        }
+        else if constexpr(F == GELU){
+            const T a = math::FRAC_2_SQRTPI<T> * math::SQRT1_2<T> * (T)0.5;
+            return (T)0.5 * (x + x * math::tanh(DEVICE(), a * ((T)0.044715f * x * x * x + x)));
+        }
+        else if constexpr(F == TANH){
+            return math::tanh(DEVICE(), x);
+        }
+        else if constexpr(F == FAST_TANH){
+            return math::fast_tanh(DEVICE(), x);
+        }
+        else if constexpr(F == SIGMOID){
+            return (T)1 / ((T)1 + math::exp(DEVICE(), -x));
+        }
+        else{
+            return 0;
+        }
+    }
+
+    template<typename DEVICE, typename T, nn::activation_functions::ActivationFunction F>
+    RL_TOOLS_FUNCTION_PLACEMENT inline T d_activation_d_x(T x){
+        using namespace nn::activation_functions;
+        static_assert(DEVICE::TYPE == devices::Type::math, "DEVICE is not a math device");
+        static_assert(check_activation_function<F>, "Invalid activation function");
+        if constexpr(F == IDENTITY){
+            return 1;
+        }
+        else if constexpr(F == RELU){
+            return x > (T)0 ? (T)1 : (T)0;
+        }
+        else if constexpr(F == GELU){
+            const T a = math::FRAC_2_SQRTPI<T> * math::SQRT1_2<T> * (T)0.5;
+            const T b = (T)0.044715f;
+            T tanh_term = math::tanh(DEVICE(), a * (b * x * x * x + x));
+            return (T)0.5*((T)1 + tanh_term) + (T)0.5 * x * ((T)1 - tanh_term * tanh_term) * a * ((T)3 * b * x * x + (T)1);
+        }
+        else if constexpr(F == TANH){
+            T a = math::tanh(DEVICE(), x);
+            return (T)1 - a * a;
+        }
+        else if constexpr(F == FAST_TANH){
+            T a = math::fast_tanh(DEVICE(), x);
+            return (T)1 - a * a;
+        }
+        else if constexpr(F == SIGMOID){
+            return activation<DEVICE, T, SIGMOID>(x) * (1 - activation<DEVICE, T, SIGMOID>(x));
+        }
+        else{
+            return 0;
+        }
+    }
+    template<typename DEVICE, typename T, nn::activation_functions::ActivationFunction F>
+    RL_TOOLS_FUNCTION_PLACEMENT auto name(){
+        using namespace nn::activation_functions;
+        if constexpr(F == IDENTITY){
+            return "IDENTITY";
+        }
+        else if constexpr(F == RELU){
+            return "RELU";
+        }
+        else if constexpr(F == GELU){
+            return "GELU";
+        }
+        else if constexpr(F == TANH){
+            return "TANH";
+        }
+        else if constexpr(F == FAST_TANH){
+            return "FAST_TANH";
+        }
+        else if constexpr(F == SIGMOID){
+            return "SIGMOID";
+        }
+        else{
+            return "NIL";
+        }
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+
+
+#endif

hyperdrone/_vendor/rl-tools/include/rl_tools/nn/capability/capability.h

@@ -0,0 +1,61 @@
+#include "../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_NN_CAPABILITY_CAPABILITY_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_NN_CAPABILITY_CAPABILITY_H
+#include "../../rl_tools.h"
+
+#include "../parameters/parameters.h"
+
+
+// The Capability Configuraiton carries properties that apply to multiple parts of a model (e.g. the high-level mode, the layers, etc.)
+// This allows e.g. to simply change the capabilities of a model without adapting all the involved specifications individually
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools::nn{
+    enum class LayerCapability{
+        Forward, // just forward
+        Backward, // forward + backward wrt to the input
+        Gradient // forward + backward wrt to the input + backward wrt to the weights
+    };
+    /*
+     Tags that can also carry the parameter type.
+     The LayerCapability determines some fields in the layers (e.g. the intermediate outputs are required for gradient calculations, while for the backward pass wrt. the input only the pre-activations might be needed)
+     The PARAMETER_TYPE determines the type of the parameters (e.g. adam requries first and second order moments in addition to the gradient)
+     This Capability system allows the switching of models for e.g. checkpointing: We are training a full model with gradients, and optimizers state then convert it to a forward only model (just the parameters) and save it as a checkpoint.
+    */
+    namespace capability{
+        template <bool T_DYNAMIC_ALLOCATION=true, bool T_CONST=false>
+        struct Forward{
+            static constexpr LayerCapability TAG = LayerCapability::Forward;
+            using PARAMETER_TYPE = nn::parameters::Plain;
+            static constexpr auto BATCH_SIZE = 1;
+            static constexpr bool DYNAMIC_ALLOCATION = T_DYNAMIC_ALLOCATION;
+            static constexpr bool CONST = T_CONST;
+            template <bool TT_DYNAMIC_ALLOCATION=true, bool TT_CONST=false>
+            using CHANGE_PARAMETERS = Forward<TT_DYNAMIC_ALLOCATION, TT_CONST>;
+        };
+        template <bool T_DYNAMIC_ALLOCATION=true, bool T_CONST=false>
+        struct Backward{
+            static constexpr LayerCapability TAG = LayerCapability::Backward;
+            using PARAMETER_TYPE = nn::parameters::Plain;
+            static constexpr bool DYNAMIC_ALLOCATION = T_DYNAMIC_ALLOCATION;
+            static constexpr bool CONST = T_CONST;
+            template <bool TT_DYNAMIC_ALLOCATION=true, bool TT_CONST=false>
+            using CHANGE_PARAMETERS = Backward<TT_DYNAMIC_ALLOCATION, TT_CONST>;
+        };
+        template <typename T_PARAMETER_TYPE, bool T_DYNAMIC_ALLOCATION=true, bool T_CONST=false>
+        struct Gradient{
+            static constexpr LayerCapability TAG = LayerCapability::Gradient;
+            using PARAMETER_TYPE = T_PARAMETER_TYPE;
+            static constexpr bool DYNAMIC_ALLOCATION = T_DYNAMIC_ALLOCATION;
+            static_assert(!utils::typing::is_same_v<T_PARAMETER_TYPE, nn::parameters::Plain>);
+            static constexpr bool CONST = T_CONST;
+            template <bool TT_DYNAMIC_ALLOCATION=true, bool TT_CONST=false>
+            using CHANGE_PARAMETERS = Gradient<PARAMETER_TYPE, TT_DYNAMIC_ALLOCATION, TT_CONST>;
+        };
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+
+#endif
+

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

@@ -0,0 +1,31 @@
+#include "../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_NN_CAPABILITY_PERSIST_CODE_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_NN_CAPABILITY_PERSIST_CODE_H
+#include "../../rl_tools.h"
+
+#include "../../nn_models/group/persist_code.h"
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools{
+    template <bool DYNAMIC_ALLOCATION, bool CONST>
+    std::string to_string(nn::capability::Forward<DYNAMIC_ALLOCATION, CONST>){
+        return std::string("RL_TOOLS""_NAMESPACE_WRAPPER ::rl_tools::nn::capability::Forward<") + (DYNAMIC_ALLOCATION ? "true" : "false") + ", " + (CONST ? "true" : "false") + ">";
+    }
+    template <bool DYNAMIC_ALLOCATION, bool CONST>
+    std::string to_string(nn::capability::Backward<DYNAMIC_ALLOCATION, CONST>){
+        return "RL_TOOLS""_NAMESPACE_WRAPPER ::rl_tools::nn::capability::Backward<" + std::string((DYNAMIC_ALLOCATION ? "true" : "false")) + ", " + (CONST ? "true" : "false") + ">";
+    }
+    template <typename T_PARAMETER_TYPE, bool DYNAMIC_ALLOCATION, bool CONST>
+    std::string to_string(nn::capability::Gradient<T_PARAMETER_TYPE, DYNAMIC_ALLOCATION, CONST>){
+        return "RL_TOOLS""_NAMESPACE_WRAPPER ::rl_tools::nn::capability::Gradient<"+ get_type_string(T_PARAMETER_TYPE{}) + std::string(", ") + std::string((DYNAMIC_ALLOCATION ? "true" : "false")) + ", " + (CONST ? "true" : "false") + ">";
+    }
+    template<typename CAPABILITY, typename TAG>
+    std::string to_string(nn::group::Capability<CAPABILITY, TAG>){
+        return std::string("RL_TOOLS_NAMESPACE_WRAPPER ::rl_tools::nn::group::Capability<") + to_string(CAPABILITY{}) + ", " + nn::group::CodeName<TAG>::VALUE + ">";
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+
+
+#endif

hyperdrone/_vendor/rl-tools/include/rl_tools/nn/group.h

@@ -0,0 +1,99 @@
+#include "../version.h"
+#include "../rl_tools.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_NN_GROUP_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_NN_GROUP_H
+
+#include "../utils/generic/typing.h"
+#include "../numeric_types/categories.h"
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools::nn::group{
+    struct Default{};
+
+    template<typename... TAGS>
+    struct Membership{
+        template<typename TAG>
+        static constexpr bool CONTAINS = (utils::typing::is_same_v<TAG, TAGS> || ...);
+    };
+
+    template<typename T, typename = void>
+    struct MembershipOf{
+        using TYPE = Membership<Default>;
+    };
+    template<typename T>
+    struct MembershipOf<T, utils::typing::void_t<typename T::MODEL_GROUPS>>{
+        using TYPE = typename T::MODEL_GROUPS;
+    };
+
+    template<typename MEMBERSHIP, typename TAG>
+    struct Append;
+    template<typename... TAGS, typename TAG>
+    struct Append<Membership<TAGS...>, TAG>{
+        using TYPE = Membership<TAGS..., TAG>;
+    };
+    template<typename TAG>
+    struct Append<Membership<Default>, TAG>{
+        using TYPE = Membership<TAG>;
+    };
+
+    template<typename... TAGS>
+    struct Only{
+        template<typename TARGET>
+        static constexpr bool MATCH = (MembershipOf<TARGET>::TYPE::template CONTAINS<TAGS> || ...);
+    };
+    struct All{
+        template<typename>
+        static constexpr bool MATCH = true;
+    };
+    template<typename TAG>
+    struct ParameterGroup{
+        template<typename TARGET>
+        static constexpr bool MATCH = utils::typing::is_same_v<typename TARGET::GROUP_TAG, TAG>;
+    };
+    template<typename TAG>
+    struct Category{
+        template<typename TARGET>
+        static constexpr bool MATCH = utils::typing::is_same_v<typename TARGET::CATEGORY_TAG, TAG>;
+    };
+    template<typename... SELECTORS>
+    struct Any{
+        template<typename TARGET>
+        static constexpr bool MATCH = (SELECTORS::template MATCH<TARGET> || ...);
+    };
+    template<typename... SELECTORS>
+    struct Every{
+        template<typename TARGET>
+        static constexpr bool MATCH = (SELECTORS::template MATCH<TARGET> && ...);
+    };
+    template<typename SELECTOR>
+    struct Not{
+        template<typename TARGET>
+        static constexpr bool MATCH = !SELECTOR::template MATCH<TARGET>;
+    };
+
+    template<typename PARAMETER_TYPE, typename GROUPS>
+    struct Parameters: PARAMETER_TYPE{
+        template<typename TYPE_POLICY, typename TI, typename SHAPE, typename GROUP_TAG, typename CATEGORY_TAG, bool DYNAMIC_ALLOCATION, bool CONST = false, typename NUMERIC_CATEGORY = numeric_types::categories::Parameter>
+        struct Specification: PARAMETER_TYPE::template Specification<TYPE_POLICY, TI, SHAPE, GROUP_TAG, CATEGORY_TAG, DYNAMIC_ALLOCATION, CONST, NUMERIC_CATEGORY>{
+            using MODEL_GROUPS = GROUPS;
+        };
+    };
+    template<typename PARAMETER_TYPE>
+    struct BaseParameters{
+        using TYPE = PARAMETER_TYPE;
+    };
+    template<typename PARAMETER_TYPE, typename GROUPS>
+    struct BaseParameters<Parameters<PARAMETER_TYPE, GROUPS>>: BaseParameters<PARAMETER_TYPE>{};
+    template<typename T_CAPABILITY, typename TAG>
+    struct Capability: T_CAPABILITY{
+        using MODEL_GROUPS = typename Append<typename MembershipOf<T_CAPABILITY>::TYPE, TAG>::TYPE;
+        using PARAMETER_TYPE = Parameters<typename BaseParameters<typename T_CAPABILITY::PARAMETER_TYPE>::TYPE, MODEL_GROUPS>;
+        template<bool DYNAMIC_ALLOCATION = true, bool CONST = false>
+        using CHANGE_PARAMETERS = Capability<typename T_CAPABILITY::template CHANGE_PARAMETERS<DYNAMIC_ALLOCATION, CONST>, TAG>;
+    };
+    template<typename T_CAPABILITY>
+    struct Capability<T_CAPABILITY, Default>: T_CAPABILITY{};
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+#endif

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

@@ -0,0 +1,122 @@
+#include "../../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_NN_LAYERS_AVG_POOL2D_LAYER_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_NN_LAYERS_AVG_POOL2D_LAYER_H
+#include "../../../utils/generic/typing.h"
+#include "../../../containers/tensor/tensor.h"
+#include "../../../nn/capability/capability.h"
+#include "../../../nn/parameters/parameters.h"
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools::nn::layers::avg_pool2d {
+
+    template<typename T_TYPE_POLICY, typename T_TI>
+    struct Configuration{
+        using TYPE_POLICY = T_TYPE_POLICY;
+        using TI = T_TI;
+    };
+
+    template <typename T_CONFIG, typename T_CAPABILITY, typename T_INPUT_SHAPE>
+    struct Specification: T_CAPABILITY, T_CONFIG{
+        using CONFIG = T_CONFIG;
+        using TYPE_POLICY = typename CONFIG::TYPE_POLICY;
+        using TI = typename CONFIG::TI;
+        using CAPABILITY = T_CAPABILITY;
+        using INPUT_SHAPE = T_INPUT_SHAPE;
+        static_assert(length(INPUT_SHAPE{}) >= 4, "AvgPool2d input shape must have at least 4 dimensions (...BATCH x H x W x C)");
+        static constexpr TI INPUT_HEIGHT = get<length(INPUT_SHAPE{})-3>(INPUT_SHAPE{});
+        static constexpr TI INPUT_WIDTH = get<length(INPUT_SHAPE{})-2>(INPUT_SHAPE{});
+        static constexpr TI INPUT_CHANNELS = get_last(INPUT_SHAPE{});
+        static constexpr TI OUTPUT_CHANNELS = INPUT_CHANNELS;
+
+        using BATCH_SHAPE = tensor::PopBack<tensor::PopBack<tensor::PopBack<INPUT_SHAPE>>>;
+        // Compute INTERNAL_BATCH_SIZE = product of all dims except the last 3 (H, W, C)
+        static constexpr TI INTERNAL_BATCH_SIZE = get<0>(tensor::CumulativeProduct<BATCH_SHAPE>{});
+
+        // Output shape: (...BATCH, C)  [global average pool collapses H and W, preserves batch dims]
+        template <typename NEW_INPUT_SHAPE>
+        struct OUTPUT_SHAPE_FACTORY{
+            static_assert(length(NEW_INPUT_SHAPE{}) >= 4);
+            static constexpr TI NEW_C = get_last(NEW_INPUT_SHAPE{});
+            static_assert(NEW_C == INPUT_CHANNELS);
+            // Pop C, W, H from the back to get batch dims, then append C.
+            using NEW_BATCH_SHAPE = tensor::PopBack<tensor::PopBack<tensor::PopBack<NEW_INPUT_SHAPE>>>;
+            using SHAPE = tensor::Append<NEW_BATCH_SHAPE, NEW_C>;
+        };
+        using OUTPUT_SHAPE = typename OUTPUT_SHAPE_FACTORY<INPUT_SHAPE>::SHAPE;
+        static constexpr TI NUM_WEIGHTS = 0;
+    };
+
+    template<typename SPEC_1, typename SPEC_2>
+    constexpr bool check_spec_memory =
+        SPEC_1::INPUT_HEIGHT == SPEC_2::INPUT_HEIGHT
+        && SPEC_1::INPUT_WIDTH == SPEC_2::INPUT_WIDTH
+        && SPEC_1::INPUT_CHANNELS == SPEC_2::INPUT_CHANNELS;
+
+    template<typename SPEC_1, typename SPEC_2>
+    constexpr bool check_spec = check_spec_memory<SPEC_1, SPEC_2>;
+
+    template <typename LAYER_SPEC, typename INPUT_SPEC, typename OUTPUT_SPEC>
+    constexpr bool check_input_output =
+        length(typename INPUT_SPEC::SHAPE{}) >= 4 &&
+        length(typename OUTPUT_SPEC::SHAPE{}) >= 2 &&
+        get<length(typename INPUT_SPEC::SHAPE{})-1>(typename INPUT_SPEC::SHAPE{}) == LAYER_SPEC::INPUT_CHANNELS &&
+        get<length(typename INPUT_SPEC::SHAPE{})-2>(typename INPUT_SPEC::SHAPE{}) == LAYER_SPEC::INPUT_WIDTH &&
+        get<length(typename INPUT_SPEC::SHAPE{})-3>(typename INPUT_SPEC::SHAPE{}) == LAYER_SPEC::INPUT_HEIGHT &&
+        get_last(typename OUTPUT_SPEC::SHAPE{}) == LAYER_SPEC::OUTPUT_CHANNELS;
+
+    struct State{};
+    struct Buffer{};
+
+    template<typename T_SPEC>
+    struct LayerForward {
+        using SPEC = T_SPEC;
+        using TYPE_POLICY = typename SPEC::TYPE_POLICY;
+        using TI = typename SPEC::TI;
+        static constexpr TI INPUT_HEIGHT = SPEC::INPUT_HEIGHT;
+        static constexpr TI INPUT_WIDTH = SPEC::INPUT_WIDTH;
+        static constexpr TI INPUT_CHANNELS = SPEC::INPUT_CHANNELS;
+        static constexpr TI OUTPUT_CHANNELS = SPEC::OUTPUT_CHANNELS;
+        static constexpr TI NUM_WEIGHTS = SPEC::NUM_WEIGHTS;
+        static constexpr TI INTERNAL_BATCH_SIZE = SPEC::INTERNAL_BATCH_SIZE;
+        using INPUT_SHAPE = typename SPEC::INPUT_SHAPE;
+        template <typename NEW_INPUT_SHAPE>
+        using OUTPUT_SHAPE_FACTORY = typename SPEC::template OUTPUT_SHAPE_FACTORY<NEW_INPUT_SHAPE>::SHAPE;
+        using OUTPUT_SHAPE = typename SPEC::OUTPUT_SHAPE;
+        template<bool DYNAMIC_ALLOCATION=true>
+        using Buffer = avg_pool2d::Buffer;
+        template<bool DYNAMIC_ALLOCATION=true>
+        using State = avg_pool2d::State;
+    };
+
+    template<typename SPEC>
+    struct LayerBackward: public LayerForward<SPEC>{};
+
+    template<typename SPEC>
+    struct LayerGradient: public LayerBackward<SPEC>{
+        using T = typename SPEC::TYPE_POLICY::template GET<numeric_types::categories::Activation>;
+        using TI = typename SPEC::TI;
+        // Internal storage: [INTERNAL_BATCH_SIZE, CHANNELS] (mapped to OUTPUT_SHAPE via view_memory)
+        using OUTPUT_CONTAINER_SHAPE = tensor::Shape<TI, SPEC::INTERNAL_BATCH_SIZE, SPEC::OUTPUT_CHANNELS>;
+        using OUTPUT_CONTAINER_SPEC = tensor::Specification<T, TI, OUTPUT_CONTAINER_SHAPE, SPEC::DYNAMIC_ALLOCATION, tensor::RowMajorStride<OUTPUT_CONTAINER_SHAPE>, SPEC::CONST>;
+        using OUTPUT_CONTAINER_TYPE = Tensor<OUTPUT_CONTAINER_SPEC>;
+        OUTPUT_CONTAINER_TYPE output;
+    };
+
+    template<typename CONFIG, typename CAPABILITY, typename INPUT_SHAPE>
+    using Layer =
+        typename utils::typing::conditional_t<CAPABILITY::TAG == nn::LayerCapability::Forward,
+            LayerForward<Specification<CONFIG, CAPABILITY, INPUT_SHAPE>>,
+        typename utils::typing::conditional_t<CAPABILITY::TAG == nn::LayerCapability::Backward,
+            LayerBackward<Specification<CONFIG, CAPABILITY, INPUT_SHAPE>>,
+        typename utils::typing::conditional_t<CAPABILITY::TAG == nn::LayerCapability::Gradient,
+            LayerGradient<Specification<CONFIG, CAPABILITY, INPUT_SHAPE>>, void>>>;
+
+    template <typename CONFIG>
+    struct BindConfiguration{
+        template <typename CAPABILITY, typename INPUT_SHAPE>
+        using Layer = nn::layers::avg_pool2d::Layer<CONFIG, CAPABILITY, INPUT_SHAPE>;
+    };
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+#endif

hyperdrone/_vendor/rl-tools/include/rl_tools/nn/layers/avg_pool2d/operations_cuda.h

@@ -0,0 +1,64 @@
+#include "../../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_NN_LAYERS_AVG_POOL2D_OPERATIONS_CUDA_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_NN_LAYERS_AVG_POOL2D_OPERATIONS_CUDA_H
+#include "../../../devices/cuda.h"
+#include "../../../nn/nn.h"
+#include "../../../mode/mode.h"
+#include "layer.h"
+#include <cudnn.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<>>
+    RL_TOOLS_FUNCTION_PLACEMENT void evaluate(devices::CUDA<DEV_SPEC>& device, const nn::layers::avg_pool2d::LayerForward<LAYER_SPEC>& layer, const Tensor<INPUT_SPEC>& input, Tensor<OUTPUT_SPEC>& output, nn::layers::avg_pool2d::Buffer&, RNG& rng, const Mode<MODE>& mode = Mode<mode::Default<>>{}){
+        static_assert(nn::layers::avg_pool2d::check_input_output<LAYER_SPEC, INPUT_SPEC, OUTPUT_SPEC>);
+        using T = typename OUTPUT_SPEC::T;
+        using TI = typename devices::CUDA<DEV_SPEC>::index_t;
+        constexpr TI N = LAYER_SPEC::INTERNAL_BATCH_SIZE, IH = LAYER_SPEC::INPUT_HEIGHT, IW = LAYER_SPEC::INPUT_WIDTH, C = LAYER_SPEC::INPUT_CHANNELS;
+        constexpr cudnnDataType_t dt = nn::cuda::get_cudnn_dtype<T>();
+        cudnnTensorDescriptor_t xd, yd;
+        check_cudnn_call(device, cudnnCreateTensorDescriptor(&xd), "cudnnCreateTensorDescriptor avgpool.xd");
+        check_cudnn_call(device, cudnnSetTensor4dDescriptor(xd, CUDNN_TENSOR_NHWC, dt, N, C, IH, IW), "cudnnSetTensor4dDescriptor avgpool.xd");
+        check_cudnn_call(device, cudnnCreateTensorDescriptor(&yd), "cudnnCreateTensorDescriptor avgpool.yd");
+        check_cudnn_call(device, cudnnSetTensor4dDescriptor(yd, CUDNN_TENSOR_NHWC, dt, N, C, 1, 1), "cudnnSetTensor4dDescriptor avgpool.yd");
+        cudnnPoolingDescriptor_t pd;
+        check_cudnn_call(device, cudnnCreatePoolingDescriptor(&pd), "cudnnCreatePoolingDescriptor avgpool.pd");
+        check_cudnn_call(device, cudnnSetPooling2dDescriptor(pd, CUDNN_POOLING_AVERAGE_COUNT_INCLUDE_PADDING, CUDNN_NOT_PROPAGATE_NAN, IH, IW, 0, 0, 1, 1), "cudnnSetPooling2dDescriptor avgpool.pd");
+        float a = 1, b = 0;
+        check_cudnn_call(device, cudnnPoolingForward(device.cudnn_handle, pd, &a, xd, input._data, &b, yd, output._data), "cudnnPoolingForward avgpool");
+        check_cudnn_call(device, cudnnDestroyPoolingDescriptor(pd), "cudnnDestroyPoolingDescriptor avgpool.pd");
+        check_cudnn_call(device, cudnnDestroyTensorDescriptor(xd), "cudnnDestroyTensorDescriptor avgpool.xd");
+        check_cudnn_call(device, cudnnDestroyTensorDescriptor(yd), "cudnnDestroyTensorDescriptor avgpool.yd");
+        check_status(device);
+    }
+    template<typename DEV_SPEC, typename LAYER_SPEC, typename INPUT_SPEC, typename OUTPUT_SPEC, typename RNG, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT void forward(devices::CUDA<DEV_SPEC>& device, nn::layers::avg_pool2d::LayerBackward<LAYER_SPEC>& layer, const Tensor<INPUT_SPEC>& input, Tensor<OUTPUT_SPEC>& output, nn::layers::avg_pool2d::Buffer& buffer, RNG& rng, const Mode<MODE>& mode = Mode<mode::Default<>>{}){
+        evaluate(device, static_cast<const nn::layers::avg_pool2d::LayerForward<LAYER_SPEC>&>(layer), input, output, buffer, rng, mode);
+    }
+    template<typename DEV_SPEC, typename LAYER_SPEC, typename INPUT_SPEC, typename D_OUTPUT_SPEC, typename MODE = mode::Default<>>
+    void backward(devices::CUDA<DEV_SPEC>& device, nn::layers::avg_pool2d::LayerGradient<LAYER_SPEC>& layer, const Tensor<INPUT_SPEC>& input, Tensor<D_OUTPUT_SPEC>& d_output, nn::layers::avg_pool2d::Buffer&, const Mode<MODE>& mode = Mode<mode::Default<>>{}){
+    }
+    template<typename DEV_SPEC, 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(devices::CUDA<DEV_SPEC>& device, nn::layers::avg_pool2d::LayerGradient<LAYER_SPEC>& layer, const Tensor<INPUT_SPEC>& input, Tensor<D_OUTPUT_SPEC>& d_output, Tensor<D_INPUT_SPEC>& d_input, nn::layers::avg_pool2d::Buffer& buffer, const Mode<MODE>& mode = Mode<mode::Default<>>{}){
+        using T = typename D_OUTPUT_SPEC::T;
+        using TI = typename devices::CUDA<DEV_SPEC>::index_t;
+        constexpr TI N = LAYER_SPEC::INTERNAL_BATCH_SIZE, IH = LAYER_SPEC::INPUT_HEIGHT, IW = LAYER_SPEC::INPUT_WIDTH, C = LAYER_SPEC::INPUT_CHANNELS;
+        constexpr cudnnDataType_t dt = nn::cuda::get_cudnn_dtype<T>();
+        cudnnTensorDescriptor_t xd, yd;
+        check_cudnn_call(device, cudnnCreateTensorDescriptor(&xd), "cudnnCreateTensorDescriptor avgpool_bwd.xd");
+        check_cudnn_call(device, cudnnSetTensor4dDescriptor(xd, CUDNN_TENSOR_NHWC, dt, N, C, IH, IW), "cudnnSetTensor4dDescriptor avgpool_bwd.xd");
+        check_cudnn_call(device, cudnnCreateTensorDescriptor(&yd), "cudnnCreateTensorDescriptor avgpool_bwd.yd");
+        check_cudnn_call(device, cudnnSetTensor4dDescriptor(yd, CUDNN_TENSOR_NHWC, dt, N, C, 1, 1), "cudnnSetTensor4dDescriptor avgpool_bwd.yd");
+        cudnnPoolingDescriptor_t pd;
+        check_cudnn_call(device, cudnnCreatePoolingDescriptor(&pd), "cudnnCreatePoolingDescriptor avgpool_bwd.pd");
+        check_cudnn_call(device, cudnnSetPooling2dDescriptor(pd, CUDNN_POOLING_AVERAGE_COUNT_INCLUDE_PADDING, CUDNN_NOT_PROPAGATE_NAN, IH, IW, 0, 0, 1, 1), "cudnnSetPooling2dDescriptor avgpool_bwd.pd");
+        float a = 1, b = 0;
+        check_cudnn_call(device, cudnnPoolingBackward(device.cudnn_handle, pd, &a, yd, layer.output._data, yd, d_output._data, xd, input._data, &b, xd, d_input._data), "cudnnPoolingBackward avgpool");
+        check_cudnn_call(device, cudnnDestroyPoolingDescriptor(pd), "cudnnDestroyPoolingDescriptor avgpool_bwd.pd");
+        check_cudnn_call(device, cudnnDestroyTensorDescriptor(xd), "cudnnDestroyTensorDescriptor avgpool_bwd.xd");
+        check_cudnn_call(device, cudnnDestroyTensorDescriptor(yd), "cudnnDestroyTensorDescriptor avgpool_bwd.yd");
+        check_status(device);
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+#endif