hyperverse 0.0.1__py3-none-any.whl

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

@@ -0,0 +1,141 @@
+#include "../../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_NN_LAYERS_UPSAMPLE2D_OPERATIONS_CUDA_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_NN_LAYERS_UPSAMPLE2D_OPERATIONS_CUDA_H
+#include "../../../devices/cuda.h"
+#include "../../../nn/nn.h"
+#include "../../../mode/mode.h"
+#include "layer.h"
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools::nn::layers::upsample2d::cuda {
+    template<typename T, typename TI>
+    __global__ void bilinear_upsample_forward(
+        const T* __restrict__ input, T* __restrict__ output,
+        TI N, TI IH, TI IW, TI OH, TI OW, TI C
+    ) {
+        TI idx = (TI)blockIdx.x * blockDim.x + threadIdx.x;
+        TI total = N * OH * OW * C;
+        if (idx >= total) return;
+        TI c_i = idx % C;
+        TI ox = (idx / C) % OW;
+        TI oy = (idx / (C * OW)) % OH;
+        TI bi = idx / (C * OW * OH);
+        float src_y = ((float)oy + 0.5f) * (float)IH / (float)OH - 0.5f;
+        float src_x = ((float)ox + 0.5f) * (float)IW / (float)OW - 0.5f;
+        int iy0 = src_y < 0.f ? 0 : (int)src_y;
+        int ix0 = src_x < 0.f ? 0 : (int)src_x;
+        int iy1 = iy0 + 1 < (int)IH ? iy0 + 1 : (int)IH - 1;
+        int ix1 = ix0 + 1 < (int)IW ? ix0 + 1 : (int)IW - 1;
+        if (iy0 >= (int)IH) iy0 = (int)IH - 1;
+        if (ix0 >= (int)IW) ix0 = (int)IW - 1;
+        float fy = src_y - (float)iy0;
+        float fx = src_x - (float)ix0;
+        if (fy < 0.f) fy = 0.f;
+        if (fx < 0.f) fx = 0.f;
+        const T* base = input + (size_t)bi * IH * IW * C;
+        float v = (1.f - fy) * (1.f - fx) * (float)base[((size_t)iy0 * IW + ix0) * C + c_i]
+                + (1.f - fy) * fx          * (float)base[((size_t)iy0 * IW + ix1) * C + c_i]
+                + fy          * (1.f - fx) * (float)base[((size_t)iy1 * IW + ix0) * C + c_i]
+                + fy          * fx          * (float)base[((size_t)iy1 * IW + ix1) * C + c_i];
+        output[idx] = (T)v;
+    }
+
+    // Input-centric backward: one thread per input element, gathers from output pixels
+    // that reference it. Deterministic — no atomicAdd needed.
+    template<typename T, typename TI>
+    __global__ void bilinear_upsample_backward(
+        const T* __restrict__ d_output, T* __restrict__ d_input,
+        TI N, TI IH, TI IW, TI OH, TI OW, TI C
+    ) {
+        TI idx = (TI)blockIdx.x * blockDim.x + threadIdx.x;
+        TI total = N * IH * IW * C;
+        if (idx >= total) return;
+        TI c_i = idx % C;
+        TI ix = (idx / C) % IW;
+        TI iy = (idx / (C * IW)) % IH;
+        TI bi = idx / (C * IW * IH);
+        float oy_lo = ((float)iy - 0.5f) * (float)OH / (float)IH - 0.5f;
+        float oy_hi = ((float)iy + 1.5f) * (float)OH / (float)IH - 0.5f;
+        int oy_start = 0 > (int)ceilf(oy_lo) ? 0 : (int)ceilf(oy_lo);
+        int oy_end = (int)OH < (int)ceilf(oy_hi) ? (int)OH : (int)ceilf(oy_hi);
+        float ox_lo = ((float)ix - 0.5f) * (float)OW / (float)IW - 0.5f;
+        float ox_hi = ((float)ix + 1.5f) * (float)OW / (float)IW - 0.5f;
+        int ox_start = 0 > (int)ceilf(ox_lo) ? 0 : (int)ceilf(ox_lo);
+        int ox_end = (int)OW < (int)ceilf(ox_hi) ? (int)OW : (int)ceilf(ox_hi);
+        float acc = 0.f;
+        const T* d_out_base = d_output + (size_t)bi * OH * OW * C;
+        for (int oy = oy_start; oy < oy_end; oy++) {
+            float src_y = ((float)oy + 0.5f) * (float)IH / (float)OH - 0.5f;
+            int iy0 = src_y < 0.f ? 0 : (int)src_y;
+            int iy1 = iy0 + 1 < (int)IH ? iy0 + 1 : (int)IH - 1;
+            if (iy0 >= (int)IH) iy0 = (int)IH - 1;
+            float fy = src_y - (float)iy0;
+            if (fy < 0.f) fy = 0.f;
+            float wy;
+            if ((int)iy == iy0 && (int)iy == iy1) wy = 1.f;
+            else if ((int)iy == iy0) wy = 1.f - fy;
+            else if ((int)iy == iy1) wy = fy;
+            else continue;
+            for (int ox = ox_start; ox < ox_end; ox++) {
+                float src_x = ((float)ox + 0.5f) * (float)IW / (float)OW - 0.5f;
+                int ix0 = src_x < 0.f ? 0 : (int)src_x;
+                int ix1 = ix0 + 1 < (int)IW ? ix0 + 1 : (int)IW - 1;
+                if (ix0 >= (int)IW) ix0 = (int)IW - 1;
+                float fx = src_x - (float)ix0;
+                if (fx < 0.f) fx = 0.f;
+                float wx;
+                if ((int)ix == ix0 && (int)ix == ix1) wx = 1.f;
+                else if ((int)ix == ix0) wx = 1.f - fx;
+                else if ((int)ix == ix1) wx = fx;
+                else continue;
+                acc += wy * wx * (float)d_out_base[((size_t)oy * OW + ox) * C + c_i];
+            }
+        }
+        d_input[idx] = (T)acc;
+    }
+}
+
+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::CUDA<DEV_SPEC>& device, const nn::layers::upsample2d::LayerForward<LAYER_SPEC>& layer, const Tensor<INPUT_SPEC>& input, Tensor<OUTPUT_SPEC>& output, nn::layers::upsample2d::Buffer&, RNG& rng, const Mode<MODE>& mode = Mode<mode::Default<>>{}){
+        static_assert(nn::layers::upsample2d::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;
+        constexpr TI IH = LAYER_SPEC::INPUT_HEIGHT, IW = LAYER_SPEC::INPUT_WIDTH, C = LAYER_SPEC::INPUT_CHANNELS;
+        constexpr TI OH = LAYER_SPEC::OUTPUT_HEIGHT, OW = LAYER_SPEC::OUTPUT_WIDTH;
+        constexpr TI total = N * OH * OW * C;
+        constexpr TI BLOCK = 256;
+        constexpr TI GRID = (total + BLOCK - 1) / BLOCK;
+        nn::layers::upsample2d::cuda::bilinear_upsample_forward<T, TI><<<GRID, BLOCK, 0, device.stream>>>(
+            input._data, output._data, N, IH, IW, OH, OW, C);
+        check_status(device);
+    }
+
+    template<typename DEV_SPEC, typename LAYER_SPEC, typename INPUT_SPEC, typename OUTPUT_SPEC, typename RNG, typename MODE = mode::Default<>>
+    void forward(devices::CUDA<DEV_SPEC>& device, nn::layers::upsample2d::LayerBackward<LAYER_SPEC>& layer, const Tensor<INPUT_SPEC>& input, Tensor<OUTPUT_SPEC>& output, nn::layers::upsample2d::Buffer& buffer, RNG& rng, const Mode<MODE>& mode = Mode<mode::Default<>>{}){
+        evaluate(device, static_cast<const nn::layers::upsample2d::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::upsample2d::LayerGradient<LAYER_SPEC>& layer, const Tensor<INPUT_SPEC>& input, Tensor<D_OUTPUT_SPEC>& d_output, nn::layers::upsample2d::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<>>
+    void backward_full(devices::CUDA<DEV_SPEC>& device, nn::layers::upsample2d::LayerGradient<LAYER_SPEC>& layer, const Tensor<INPUT_SPEC>& input, Tensor<D_OUTPUT_SPEC>& d_output, Tensor<D_INPUT_SPEC>& d_input, nn::layers::upsample2d::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;
+        constexpr TI IH = LAYER_SPEC::INPUT_HEIGHT, IW = LAYER_SPEC::INPUT_WIDTH, C = LAYER_SPEC::INPUT_CHANNELS;
+        constexpr TI OH = LAYER_SPEC::OUTPUT_HEIGHT, OW = LAYER_SPEC::OUTPUT_WIDTH;
+        constexpr TI total = N * IH * IW * C;
+        constexpr TI BLOCK = 256;
+        constexpr TI GRID = (total + BLOCK - 1) / BLOCK;
+        nn::layers::upsample2d::cuda::bilinear_upsample_backward<T, TI><<<GRID, BLOCK, 0, device.stream>>>(
+            d_output._data, d_input._data, N, IH, IW, OH, OW, C);
+        check_status(device);
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+#endif

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

@@ -0,0 +1,230 @@
+#include "../../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_NN_LAYERS_UPSAMPLE2D_OPERATIONS_GENERIC_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_NN_LAYERS_UPSAMPLE2D_OPERATIONS_GENERIC_H
+#include "layer.h"
+#include "../../../rl_tools.h"
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools{
+    // ======================== malloc / free ========================
+    template<typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void malloc(DEVICE& device, nn::layers::upsample2d::LayerForward<SPEC>& layer) {}
+    template<typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void free(DEVICE& device, nn::layers::upsample2d::LayerForward<SPEC>& layer) {}
+    template<typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void malloc(DEVICE& device, nn::layers::upsample2d::LayerBackward<SPEC>& layer) {}
+    template<typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void free(DEVICE& device, nn::layers::upsample2d::LayerBackward<SPEC>& layer) {}
+    template<typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void malloc(DEVICE& device, nn::layers::upsample2d::LayerGradient<SPEC>& layer) {
+        malloc(device, layer.output);
+    }
+    template<typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void free(DEVICE& device, nn::layers::upsample2d::LayerGradient<SPEC>& layer) {
+        free(device, layer.output);
+    }
+    template<typename DEVICE>
+    RL_TOOLS_FUNCTION_PLACEMENT void malloc(DEVICE& device, nn::layers::upsample2d::State&) {}
+    template<typename DEVICE>
+    RL_TOOLS_FUNCTION_PLACEMENT void free(DEVICE& device, nn::layers::upsample2d::State&) {}
+    template<typename DEVICE>
+    RL_TOOLS_FUNCTION_PLACEMENT void malloc(DEVICE& device, nn::layers::upsample2d::Buffer&) {}
+    template<typename DEVICE>
+    RL_TOOLS_FUNCTION_PLACEMENT void free(DEVICE& device, nn::layers::upsample2d::Buffer&) {}
+    template<typename SOURCE_DEVICE, typename TARGET_DEVICE>
+    RL_TOOLS_FUNCTION_PLACEMENT void copy(SOURCE_DEVICE&, TARGET_DEVICE&, nn::layers::upsample2d::State&, nn::layers::upsample2d::State&) {}
+    template<typename SPEC, typename DEVICE, typename RNG, typename MODE>
+    RL_TOOLS_FUNCTION_PLACEMENT void reset(DEVICE&, const nn::layers::upsample2d::LayerForward<SPEC>&, nn::layers::upsample2d::State&, RNG&, Mode<MODE>) {}
+
+    // ======================== init_weights (no-op) ========================
+    template<typename DEVICE, typename SPEC, typename RNG>
+    RL_TOOLS_FUNCTION_PLACEMENT void init_weights(DEVICE& device, nn::layers::upsample2d::LayerForward<SPEC>& layer, RNG& rng) {}
+
+    // ======================== evaluate (bilinear upsample) ========================
+    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::upsample2d::LayerForward<LAYER_SPEC>& layer, const Tensor<INPUT_SPEC>& input, Tensor<OUTPUT_SPEC>& output, nn::layers::upsample2d::Buffer&, RNG& rng, const Mode<MODE>& mode = Mode<mode::Default<>>{}) {
+        static_assert(nn::layers::upsample2d::check_input_output<LAYER_SPEC, INPUT_SPEC, OUTPUT_SPEC>);
+        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 IH = LAYER_SPEC::INPUT_HEIGHT;
+        constexpr TI IW = LAYER_SPEC::INPUT_WIDTH;
+        constexpr TI OH = LAYER_SPEC::OUTPUT_HEIGHT;
+        constexpr TI OW = LAYER_SPEC::OUTPUT_WIDTH;
+        constexpr TI C = LAYER_SPEC::INPUT_CHANNELS;
+        constexpr TI BATCH_SIZE = product(typename INPUT_SPEC::SHAPE{}) / (IH * IW * C);
+        using INTERNAL_INPUT_SHAPE = tensor::Shape<TI, BATCH_SIZE, IH, IW, C>;
+        using INTERNAL_OUTPUT_SHAPE = tensor::Shape<TI, BATCH_SIZE, OH, OW, C>;
+        auto input_4d = view_memory<INTERNAL_INPUT_SHAPE>(device, input);
+        auto output_4d = view_memory<INTERNAL_OUTPUT_SHAPE>(device, output);
+        for(TI bi = 0; bi < BATCH_SIZE; bi++){
+            for(TI oy = 0; oy < OH; oy++){
+                for(TI ox = 0; ox < OW; ox++){
+                    // Source coordinate (align_corners=false, half-pixel offset)
+                    ACCUMULATOR_TYPE sy = ((ACCUMULATOR_TYPE)oy + (ACCUMULATOR_TYPE)0.5) * (ACCUMULATOR_TYPE)IH / (ACCUMULATOR_TYPE)OH - (ACCUMULATOR_TYPE)0.5;
+                    ACCUMULATOR_TYPE sx = ((ACCUMULATOR_TYPE)ox + (ACCUMULATOR_TYPE)0.5) * (ACCUMULATOR_TYPE)IW / (ACCUMULATOR_TYPE)OW - (ACCUMULATOR_TYPE)0.5;
+                    TI y0 = sy < 0 ? 0 : (TI)sy;
+                    TI x0 = sx < 0 ? 0 : (TI)sx;
+                    TI y1 = y0 + 1 < IH ? y0 + 1 : IH - 1;
+                    TI x1 = x0 + 1 < IW ? x0 + 1 : IW - 1;
+                    if(y0 >= IH) y0 = IH - 1;
+                    if(x0 >= IW) x0 = IW - 1;
+                    ACCUMULATOR_TYPE fy = sy - (ACCUMULATOR_TYPE)y0;
+                    ACCUMULATOR_TYPE fx = sx - (ACCUMULATOR_TYPE)x0;
+                    if(fy < 0) fy = 0;
+                    if(fx < 0) fx = 0;
+                    ACCUMULATOR_TYPE w00 = ((ACCUMULATOR_TYPE)1 - fy) * ((ACCUMULATOR_TYPE)1 - fx);
+                    ACCUMULATOR_TYPE w01 = ((ACCUMULATOR_TYPE)1 - fy) * fx;
+                    ACCUMULATOR_TYPE w10 = fy * ((ACCUMULATOR_TYPE)1 - fx);
+                    ACCUMULATOR_TYPE w11 = fy * fx;
+                    for(TI c_i = 0; c_i < C; c_i++){
+                        ACCUMULATOR_TYPE v = w00 * (ACCUMULATOR_TYPE)get(device, input_4d, bi, y0, x0, c_i)
+                                           + w01 * (ACCUMULATOR_TYPE)get(device, input_4d, bi, y0, x1, c_i)
+                                           + w10 * (ACCUMULATOR_TYPE)get(device, input_4d, bi, y1, x0, c_i)
+                                           + w11 * (ACCUMULATOR_TYPE)get(device, input_4d, bi, y1, x1, c_i);
+                        set(device, output_4d, (T)v, bi, oy, ox, c_i);
+                    }
+                }
+            }
+        }
+    }
+
+    template<typename DEVICE, typename LAYER_SPEC, typename INPUT_SPEC, typename OUTPUT_SPEC, typename RNG, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT void evaluate_step(DEVICE& device, const nn::layers::upsample2d::LayerForward<LAYER_SPEC>& layer, const Tensor<INPUT_SPEC>& input, nn::layers::upsample2d::State& state, Tensor<OUTPUT_SPEC>& output, nn::layers::upsample2d::Buffer& buffer, RNG& rng, const Mode<MODE>& mode = Mode<mode::Default<>>{}){
+        evaluate(device, layer, input, 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::upsample2d::LayerBackward<LAYER_SPEC>& layer, const Tensor<INPUT_SPEC>& input, Tensor<OUTPUT_SPEC>& output, nn::layers::upsample2d::Buffer& buffer, RNG& rng, const Mode<MODE>& mode = Mode<mode::Default<>>{}){
+        evaluate(device, static_cast<const nn::layers::upsample2d::LayerForward<LAYER_SPEC>&>(layer), input, output, buffer, rng, mode);
+    }
+    template<typename DEVICE, typename LAYER_SPEC, typename INPUT_SPEC, typename RNG, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT void forward(DEVICE& device, nn::layers::upsample2d::LayerGradient<LAYER_SPEC>& layer, const Tensor<INPUT_SPEC>& input, nn::layers::upsample2d::Buffer& buffer, RNG& rng, const Mode<MODE>& mode = Mode<mode::Default<>>{}) {
+        forward(device, static_cast<nn::layers::upsample2d::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::upsample2d::LayerGradient<LAYER_SPEC>& layer, const Tensor<INPUT_SPEC>& input, Tensor<OUTPUT_SPEC>& output, nn::layers::upsample2d::Buffer& buffer, RNG& rng, const Mode<MODE>& mode = Mode<mode::Default<>>{}) {
+        forward(device, layer, input, buffer, rng, mode);
+        copy(device, device, layer.output, output);
+    }
+
+    // ======================== 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::upsample2d::LayerBackward<LAYER_SPEC>& layer, const Tensor<D_OUTPUT_SPEC>& d_output, Tensor<D_INPUT_SPEC>& d_input, nn::layers::upsample2d::Buffer&, const Mode<MODE>& mode = Mode<mode::Default<>>{}){
+        static_assert(nn::layers::upsample2d::check_input_output<LAYER_SPEC, D_INPUT_SPEC, D_OUTPUT_SPEC>);
+        using TI = typename DEVICE::index_t;
+        using T = typename D_OUTPUT_SPEC::T;
+        using ACCUMULATOR_TYPE = typename LAYER_SPEC::TYPE_POLICY::template GET<numeric_types::categories::Accumulator>;
+        constexpr TI IH = LAYER_SPEC::INPUT_HEIGHT;
+        constexpr TI IW = LAYER_SPEC::INPUT_WIDTH;
+        constexpr TI OH = LAYER_SPEC::OUTPUT_HEIGHT;
+        constexpr TI OW = LAYER_SPEC::OUTPUT_WIDTH;
+        constexpr TI C = LAYER_SPEC::INPUT_CHANNELS;
+        constexpr TI BATCH_SIZE = LAYER_SPEC::INTERNAL_BATCH_SIZE;
+        using INTERNAL_D_OUTPUT_SHAPE = tensor::Shape<TI, BATCH_SIZE, OH, OW, C>;
+        using INTERNAL_D_INPUT_SHAPE = tensor::Shape<TI, BATCH_SIZE, IH, IW, C>;
+        auto d_output_4d = view_memory<INTERNAL_D_OUTPUT_SHAPE>(device, d_output);
+        auto d_input_4d = view_memory<INTERNAL_D_INPUT_SHAPE>(device, d_input);
+        // Zero d_input first (scatter-add pattern)
+        set_all(device, d_input_4d, (T)0);
+        for(TI bi = 0; bi < BATCH_SIZE; bi++){
+            for(TI oy = 0; oy < OH; oy++){
+                for(TI ox = 0; ox < OW; ox++){
+                    ACCUMULATOR_TYPE sy = ((ACCUMULATOR_TYPE)oy + (ACCUMULATOR_TYPE)0.5) * (ACCUMULATOR_TYPE)IH / (ACCUMULATOR_TYPE)OH - (ACCUMULATOR_TYPE)0.5;
+                    ACCUMULATOR_TYPE sx = ((ACCUMULATOR_TYPE)ox + (ACCUMULATOR_TYPE)0.5) * (ACCUMULATOR_TYPE)IW / (ACCUMULATOR_TYPE)OW - (ACCUMULATOR_TYPE)0.5;
+                    TI y0 = sy < 0 ? 0 : (TI)sy;
+                    TI x0 = sx < 0 ? 0 : (TI)sx;
+                    TI y1 = y0 + 1 < IH ? y0 + 1 : IH - 1;
+                    TI x1 = x0 + 1 < IW ? x0 + 1 : IW - 1;
+                    if(y0 >= IH) y0 = IH - 1;
+                    if(x0 >= IW) x0 = IW - 1;
+                    ACCUMULATOR_TYPE fy = sy - (ACCUMULATOR_TYPE)y0;
+                    ACCUMULATOR_TYPE fx = sx - (ACCUMULATOR_TYPE)x0;
+                    if(fy < 0) fy = 0;
+                    if(fx < 0) fx = 0;
+                    ACCUMULATOR_TYPE w00 = ((ACCUMULATOR_TYPE)1 - fy) * ((ACCUMULATOR_TYPE)1 - fx);
+                    ACCUMULATOR_TYPE w01 = ((ACCUMULATOR_TYPE)1 - fy) * fx;
+                    ACCUMULATOR_TYPE w10 = fy * ((ACCUMULATOR_TYPE)1 - fx);
+                    ACCUMULATOR_TYPE w11 = fy * fx;
+                    for(TI c_i = 0; c_i < C; c_i++){
+                        ACCUMULATOR_TYPE grad = (ACCUMULATOR_TYPE)get(device, d_output_4d, bi, oy, ox, c_i);
+                        set(device, d_input_4d, (T)((ACCUMULATOR_TYPE)get(device, d_input_4d, bi, y0, x0, c_i) + w00 * grad), bi, y0, x0, c_i);
+                        set(device, d_input_4d, (T)((ACCUMULATOR_TYPE)get(device, d_input_4d, bi, y0, x1, c_i) + w01 * grad), bi, y0, x1, c_i);
+                        set(device, d_input_4d, (T)((ACCUMULATOR_TYPE)get(device, d_input_4d, bi, y1, x0, c_i) + w10 * grad), bi, y1, x0, c_i);
+                        set(device, d_input_4d, (T)((ACCUMULATOR_TYPE)get(device, d_input_4d, bi, y1, x1, c_i) + w11 * grad), bi, y1, x1, c_i);
+                    }
+                }
+            }
+        }
+    }
+    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::upsample2d::LayerGradient<LAYER_SPEC>& layer, const Tensor<INPUT_SPEC>& input, Tensor<D_OUTPUT_SPEC>& d_output, nn::layers::upsample2d::Buffer&, const Mode<MODE>& mode = Mode<mode::Default<>>{}){
+        // No parameters to accumulate gradients for
+    }
+    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::upsample2d::LayerGradient<LAYER_SPEC>& layer, const Tensor<INPUT_SPEC>& input, Tensor<D_OUTPUT_SPEC>& d_output, Tensor<D_INPUT_SPEC>& d_input, nn::layers::upsample2d::Buffer& buffer, const Mode<MODE>& mode = Mode<mode::Default<>>{}){
+        backward_input(device, static_cast<const nn::layers::upsample2d::LayerBackward<LAYER_SPEC>&>(layer), d_output, d_input, buffer, mode);
+    }
+
+    // ======================== zero_gradient / update / _reset_optimizer_state (no-ops) ========================
+    template<typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void zero_gradient(DEVICE& device, nn::layers::upsample2d::LayerGradient<SPEC>& layer) {}
+    template<typename DEVICE, typename SOURCE_SPEC, typename TARGET_SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void add_gradient(DEVICE& device, nn::layers::upsample2d::LayerGradient<SOURCE_SPEC>& source, nn::layers::upsample2d::LayerGradient<TARGET_SPEC>& target) {}
+    template<typename SOURCE_DEVICE, typename TARGET_DEVICE, typename SOURCE_SPEC, typename TARGET_SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void copy_gradient(SOURCE_DEVICE& source_device, TARGET_DEVICE& target_device, const nn::layers::upsample2d::LayerGradient<SOURCE_SPEC>& source, nn::layers::upsample2d::LayerGradient<TARGET_SPEC>& target) {}
+    template<typename DEVICE, typename SPEC, typename OPTIMIZER>
+    RL_TOOLS_FUNCTION_PLACEMENT void update(DEVICE& device, nn::layers::upsample2d::LayerGradient<SPEC>& layer, OPTIMIZER& optimizer) {}
+    template<typename DEVICE, typename SPEC, typename OPTIMIZER>
+    RL_TOOLS_FUNCTION_PLACEMENT void _reset_optimizer_state(DEVICE& device, nn::layers::upsample2d::LayerGradient<SPEC>& layer, OPTIMIZER& optimizer) {}
+
+    // ======================== copy ========================
+    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::upsample2d::LayerForward<SOURCE_SPEC>& source, nn::layers::upsample2d::LayerForward<TARGET_SPEC>& target) {
+        static_assert(nn::layers::upsample2d::check_spec_memory<SOURCE_SPEC, TARGET_SPEC>);
+    }
+    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::upsample2d::LayerBackward<SOURCE_SPEC>& source, nn::layers::upsample2d::LayerBackward<TARGET_SPEC>& target) {
+        static_assert(nn::layers::upsample2d::check_spec_memory<SOURCE_SPEC, TARGET_SPEC>);
+    }
+    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::upsample2d::LayerGradient<SOURCE_SPEC>& source, nn::layers::upsample2d::LayerGradient<TARGET_SPEC>& target) {
+        static_assert(nn::layers::upsample2d::check_spec_memory<SOURCE_SPEC, TARGET_SPEC>);
+        copy(source_device, target_device, source.output, target.output);
+    }
+
+    // ======================== abs_diff ========================
+    template<typename DEVICE, typename S1, typename S2>
+    RL_TOOLS_FUNCTION_PLACEMENT typename S1::TYPE_POLICY::DEFAULT abs_diff(DEVICE& device, const nn::layers::upsample2d::LayerForward<S1>&, const nn::layers::upsample2d::LayerForward<S2>&) { return 0; }
+    template<typename DEVICE, typename S1, typename S2>
+    RL_TOOLS_FUNCTION_PLACEMENT typename S1::TYPE_POLICY::DEFAULT abs_diff(DEVICE& device, const nn::layers::upsample2d::LayerBackward<S1>&, const nn::layers::upsample2d::LayerBackward<S2>&) { return 0; }
+    template<typename DEVICE, typename S1, typename S2>
+    RL_TOOLS_FUNCTION_PLACEMENT typename S1::TYPE_POLICY::DEFAULT abs_diff(DEVICE& device, const nn::layers::upsample2d::LayerGradient<S1>& l1, const nn::layers::upsample2d::LayerGradient<S2>& l2) { return abs_diff(device, l1.output, l2.output); }
+    template<typename DEVICE>
+    RL_TOOLS_FUNCTION_PLACEMENT auto abs_diff(DEVICE& device, const nn::layers::upsample2d::State&, const nn::layers::upsample2d::State&) { return 0; }
+
+    // ======================== reset_forward_state ========================
+    template<typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void reset_forward_state(DEVICE& device, nn::layers::upsample2d::LayerBackward<SPEC>&) {}
+    template<typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void reset_forward_state(DEVICE& device, nn::layers::upsample2d::LayerGradient<SPEC>& l) { set_all(device, l.output, 0); }
+
+    // ======================== is_nan ========================
+    template<typename DEVICE, typename SPEC, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT bool is_nan(DEVICE&, const nn::layers::upsample2d::LayerForward<SPEC>&, const Mode<MODE>& = Mode<mode::Default<>>{}) { return false; }
+    template<typename DEVICE, typename SPEC, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT bool is_nan(DEVICE&, const nn::layers::upsample2d::LayerBackward<SPEC>&, const Mode<MODE>& = Mode<mode::Default<>>{}) { return false; }
+    template<typename DEVICE, typename SPEC, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT bool is_nan(DEVICE& device, const nn::layers::upsample2d::LayerGradient<SPEC>& l, const Mode<MODE>& mode = Mode<mode::Default<>>{}) { return is_nan(device, l.output, mode); }
+    template<typename DEVICE, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT bool is_nan(DEVICE&, nn::layers::upsample2d::State&, const Mode<MODE>& = Mode<mode::Default<>>{}) { return false; }
+
+    // ======================== output ========================
+    template<typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT auto output(DEVICE& device, nn::layers::upsample2d::LayerGradient<SPEC>& l){
+        return view_memory<typename SPEC::OUTPUT_SHAPE>(device, l.output);
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+#endif

hyperdrone/_vendor/rl-tools/include/rl_tools/nn/loss_functions/categorical_cross_entropy/operations_generic.h

@@ -0,0 +1,102 @@
+#include "../../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_NN_LOSS_FUNCTIONS_CATEGORICAL_CROSS_ENTROPY_OPERATIONS_GENERIC_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_NN_LOSS_FUNCTIONS_CATEGORICAL_CROSS_ENTROPY_OPERATIONS_GENERIC_H
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools::nn::loss_functions::categorical_cross_entropy{
+    template<typename DEVICE, typename SPEC_A, typename SPEC_B>
+    RL_TOOLS_FUNCTION_PLACEMENT typename SPEC_A::T evaluate(DEVICE& device, Matrix<SPEC_A> a, Matrix<SPEC_B> b, typename SPEC_A::T loss_weight = 1) {
+        // a is logits, b is indices resembling empirical one-hot distributions
+        static_assert(SPEC_A::ROWS == SPEC_B::ROWS);
+        static_assert(SPEC_B::COLS == 1);
+        using T = typename SPEC_A::T;
+        using TI = typename SPEC_A::TI;
+        T acc = 0;
+        for(TI row_i = 0; row_i < SPEC_A::ROWS; row_i++) {
+            acc += get(a, row_i, get(b, row_i, 0));
+            T maximum = max(device, row(device, a, row_i));
+            acc -= maximum;
+            T sum = 0;
+            for(TI col_i = 0; col_i < SPEC_A::COLS; col_i++){
+                T logit = get(a, row_i, col_i);
+                sum += math::exp(device.math, logit - maximum);
+            }
+            acc -= math::log(device.math, sum);
+        }
+        return -acc * loss_weight / SPEC_A::ROWS;
+    }
+    template<typename DEVICE, typename SPEC_A, typename SPEC_B, typename SPEC_DA>
+    RL_TOOLS_FUNCTION_PLACEMENT void gradient(DEVICE& device, Matrix<SPEC_A> a, Matrix<SPEC_B> b, Matrix<SPEC_DA> d_a, typename SPEC_A::T loss_weight = 1) {
+        static_assert(SPEC_A::ROWS == SPEC_B::ROWS);
+        static_assert(SPEC_B::COLS == 1);
+        static_assert(containers::check_structure<SPEC_A, SPEC_DA>);
+        using T = typename SPEC_A::T;
+        using TI = typename SPEC_A::TI;
+        T effective_weight = loss_weight / SPEC_A::ROWS;
+        set_all(device, d_a, 0);
+        for(TI row_i = 0; row_i < SPEC_A::ROWS; row_i++) {
+            T maximum = max(device, row(device, a, row_i));
+            T sum = 0;
+            for(TI col_i = 0; col_i < SPEC_A::COLS; col_i++) {
+                T logit = get(a, row_i, col_i);
+                sum += math::exp(device.math, logit - maximum);
+            }
+            T sum_reciprocal = (T)1.0 / sum;
+            TI target_index = get(b, row_i, 0);
+            increment(d_a, row_i, target_index, -1 * effective_weight);
+            for(TI col_i = 0; col_i < SPEC_A::COLS; col_i++) {
+                T p = math::exp(device.math, get(a, row_i, col_i) - maximum) * sum_reciprocal;
+                increment(d_a, row_i, col_i, p*effective_weight);
+            }
+        }
+    }
+    template<typename DEVICE, typename SPEC_A, typename SPEC_B, typename SPEC_DA>
+    RL_TOOLS_FUNCTION_PLACEMENT void gradient_tiled(DEVICE& device, Matrix<SPEC_A> a, Matrix<SPEC_B> b, Matrix<SPEC_DA> d_a, typename SPEC_A::T loss_weight = 1) {
+        static_assert(SPEC_A::ROWS == SPEC_B::ROWS);
+        static_assert(SPEC_B::COLS == 1);
+        static_assert(containers::check_structure<SPEC_A, SPEC_DA>);
+        using T = typename SPEC_A::T;
+        using TI = typename SPEC_A::TI;
+
+        T effective_weight = loss_weight / SPEC_A::ROWS;
+        set_all(device, d_a, 0);
+        for(TI row_i = 0; row_i < SPEC_A::ROWS; row_i++) {
+            T maximum = max(device, row(device, a, row_i));
+            T sum = 0;
+            constexpr TI TILE_SIZE = 8;
+            constexpr TI REMAINDER = SPEC_A::COLS % TILE_SIZE;
+            constexpr TI REMAINDER_START = SPEC_A::COLS - REMAINDER;
+
+            for(TI tile_i = 0; tile_i < SPEC_A::COLS/TILE_SIZE; tile_i++) {
+                for(TI col_i = tile_i * TILE_SIZE; col_i < (tile_i+1)*TILE_SIZE; col_i++) {
+                    T logit = get(a, row_i, col_i);
+                    sum += math::exp(device.math, logit - maximum);
+                }
+            }
+            for(TI col_i = REMAINDER_START; col_i < SPEC_A::COLS; col_i++) {
+                T logit = get(a, row_i, col_i);
+                sum += math::exp(device.math, logit - maximum);
+            }
+
+            T sum_reciprocal = (T)1.0 / sum;
+
+            TI target_index = get(b, row_i, 0);
+            increment(d_a, row_i, target_index, -1 * effective_weight);
+            for(TI tile_i = 0; tile_i < SPEC_A::COLS/TILE_SIZE; tile_i++) {
+                for(TI col_i = tile_i * TILE_SIZE; col_i < (tile_i+1)*TILE_SIZE; col_i++) {
+                    T p = math::exp(device.math, get(a, row_i, col_i) - maximum) * sum_reciprocal;
+                    increment(d_a, row_i, col_i, p * effective_weight);
+                }
+            }
+            for(TI col_i = REMAINDER_START; col_i < SPEC_A::COLS; col_i++) {
+                T p = math::exp(device.math, get(a, row_i, col_i) - maximum) * sum_reciprocal;
+                increment(d_a, row_i, col_i, p * effective_weight);
+            }
+        }
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+
+#endif
+

hyperdrone/_vendor/rl-tools/include/rl_tools/nn/loss_functions/mse/operations_cuda.h

@@ -0,0 +1,97 @@
+#include "../../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_NN_LOSS_FUNCTIONS_MSE_OPERATIONS_CUDA_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_NN_LOSS_FUNCTIONS_MSE_OPERATIONS_CUDA_H
+
+#include "../../../devices/cuda.h"
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools::nn::loss_functions::mse {
+    namespace internal::mse{
+        template<typename DEV_SPEC, typename SPEC_A, typename SPEC_B, typename SPEC_DA, typename LOSS_WEIGHT_SPEC>
+        __global__
+        void d_mse_d_x_kernel(devices::CUDA<DEV_SPEC> device, Matrix<SPEC_A> a, Matrix<SPEC_B> b, Matrix<SPEC_DA> d_a, Tensor<LOSS_WEIGHT_SPEC> loss_weight) {
+            static_assert(containers::check_structure<SPEC_A, SPEC_B>);
+            static_assert(containers::check_structure<SPEC_A, SPEC_DA>);
+            static_assert(LOSS_WEIGHT_SPEC::SHAPE::LENGTH == 1);
+            static_assert(LOSS_WEIGHT_SPEC::SHAPE::template GET<0> == 1 || LOSS_WEIGHT_SPEC::SHAPE::template GET<0> == 0);
+            constexpr bool LOSS_WEIGHT_PROVIDED = LOSS_WEIGHT_SPEC::SHAPE::template GET<0> == 1;
+            using T = typename SPEC_A::T;
+            using TI = typename devices::CUDA<DEV_SPEC>::index_t;
+            constexpr TI BATCH_SIZE = SPEC_A::ROWS;
+            constexpr TI OUTPUT_DIM = SPEC_A::COLS;
+
+            TI output_pos_x = blockIdx.x * blockDim.x + threadIdx.x;
+            TI output_pos_y = blockIdx.y * blockDim.y + threadIdx.y;
+            T loss_weight_value;
+            if constexpr(LOSS_WEIGHT_PROVIDED) {
+                loss_weight_value = get(device, loss_weight, 0);
+            }
+            else {
+                loss_weight_value = (T)1;
+            }
+            if(output_pos_x < OUTPUT_DIM && output_pos_y < BATCH_SIZE){
+                T diff = get(a, output_pos_y, output_pos_x) - get(b, output_pos_y, output_pos_x);
+                set(d_a, output_pos_y, output_pos_x, (T)2*diff/(T)(SPEC_A::ROWS * SPEC_A::COLS) * loss_weight_value);
+            }
+        }
+    }
+
+    template<typename DEV_SPEC, typename SPEC_A, typename SPEC_B, typename SPEC_DA, typename LOSS_WEIGHT_SPEC>
+    void gradient(devices::CUDA<DEV_SPEC>& device, Matrix<SPEC_A>& a, Matrix<SPEC_B>& b, Matrix<SPEC_DA>& d_a, Tensor<LOSS_WEIGHT_SPEC>& loss_weight) {
+        static_assert(containers::check_structure<SPEC_A, SPEC_B>);
+        static_assert(containers::check_structure<SPEC_A, SPEC_DA>);
+        using DEVICE = devices::CUDA<DEV_SPEC>;
+        using TI = typename DEVICE::index_t;
+        constexpr TI BATCH_SIZE = SPEC_A::ROWS;
+        constexpr TI OUTPUT_DIM = SPEC_A::COLS;
+        constexpr TI BLOCKSIZE_ACTIVATION_BATCH = 32;
+        constexpr TI BLOCKSIZE_ACTIVATION_OUTPUT = 32;
+        constexpr TI N_BLOCKS_ACTIVATION_BATCH = RL_TOOLS_DEVICES_CUDA_CEIL(BATCH_SIZE, BLOCKSIZE_ACTIVATION_BATCH);
+        constexpr TI N_BLOCKS_ACTIVATION_OUTPUT = RL_TOOLS_DEVICES_CUDA_CEIL(OUTPUT_DIM, BLOCKSIZE_ACTIVATION_OUTPUT);
+        dim3 activation_grid(N_BLOCKS_ACTIVATION_OUTPUT, N_BLOCKS_ACTIVATION_BATCH);
+        dim3 activation_block(BLOCKSIZE_ACTIVATION_OUTPUT, BLOCKSIZE_ACTIVATION_BATCH);
+        using DEVICE = devices::CUDA<DEV_SPEC>;
+        devices::cuda::TAG<DEVICE, true> tag_device{};
+        internal::mse::d_mse_d_x_kernel<<<activation_grid, activation_block, 0, device.stream>>>(tag_device, a, b, d_a, loss_weight);
+        check_status(device);
+    }
+    namespace internal::mse{
+        template<typename DEV_SPEC, typename SPEC_A, typename SPEC_B, typename SPEC_DA>
+        __global__
+        void d_mse_d_x_scalar_kernel(devices::CUDA<DEV_SPEC> device, Matrix<SPEC_A> a, Matrix<SPEC_B> b, Matrix<SPEC_DA> d_a, typename SPEC_A::T loss_weight) {
+            using T = typename SPEC_A::T;
+            using TI = typename devices::CUDA<DEV_SPEC>::index_t;
+            constexpr TI BATCH_SIZE = SPEC_A::ROWS;
+            constexpr TI OUTPUT_DIM = SPEC_A::COLS;
+            TI output_pos_x = blockIdx.x * blockDim.x + threadIdx.x;
+            TI output_pos_y = blockIdx.y * blockDim.y + threadIdx.y;
+            if(output_pos_x < OUTPUT_DIM && output_pos_y < BATCH_SIZE){
+                T diff = get(a, output_pos_y, output_pos_x) - get(b, output_pos_y, output_pos_x);
+                set(d_a, output_pos_y, output_pos_x, (T)2*diff/(T)(SPEC_A::ROWS * SPEC_A::COLS) * loss_weight);
+            }
+        }
+    }
+    template<typename DEV_SPEC, typename SPEC_A, typename SPEC_B, typename SPEC_DA>
+    void gradient(devices::CUDA<DEV_SPEC>& device, Matrix<SPEC_A>& a, Matrix<SPEC_B>& b, Matrix<SPEC_DA>& d_a, typename SPEC_A::T loss_weight_scalar) {
+        static_assert(containers::check_structure<SPEC_A, SPEC_B>);
+        static_assert(containers::check_structure<SPEC_A, SPEC_DA>);
+        using DEVICE = devices::CUDA<DEV_SPEC>;
+        using TI = typename DEVICE::index_t;
+        constexpr TI BATCH_SIZE = SPEC_A::ROWS;
+        constexpr TI OUTPUT_DIM = SPEC_A::COLS;
+        constexpr TI BLOCKSIZE_BATCH = 32;
+        constexpr TI BLOCKSIZE_OUTPUT = 32;
+        constexpr TI N_BLOCKS_BATCH = RL_TOOLS_DEVICES_CUDA_CEIL(BATCH_SIZE, BLOCKSIZE_BATCH);
+        constexpr TI N_BLOCKS_OUTPUT = RL_TOOLS_DEVICES_CUDA_CEIL(OUTPUT_DIM, BLOCKSIZE_OUTPUT);
+        dim3 grid(N_BLOCKS_OUTPUT, N_BLOCKS_BATCH);
+        dim3 block(BLOCKSIZE_OUTPUT, BLOCKSIZE_BATCH);
+        devices::cuda::TAG<DEVICE, true> tag_device{};
+        internal::mse::d_mse_d_x_scalar_kernel<<<grid, block, 0, device.stream>>>(tag_device, a, b, d_a, loss_weight_scalar);
+        check_status(device);
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+#include "operations_generic.h"
+
+#endif