hyperdrone 0.1.0__py3-none-any.whl

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

@@ -0,0 +1,138 @@
+#include "../../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_NN_LAYERS_RESNET_BLOCK_OPERATIONS_CUDA_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_NN_LAYERS_RESNET_BLOCK_OPERATIONS_CUDA_H
+#include "../../../devices/cuda.h"
+#include "../../../nn/nn.h"
+#include "../../../mode/mode.h"
+#include "../conv2d/operations_generic.h"
+#include "../conv2d/operations_cuda.h"
+#include "layer.h"
+#include <cudnn.h>
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools{
+    namespace nn::layers::resnet_block::cuda::kernels{
+        template<typename DEV_SPEC, typename T>
+        __global__
+        void add_relu(devices::CUDA<DEV_SPEC> device, T* output, const T* shortcut, unsigned int n){
+            unsigned int i = blockIdx.x * blockDim.x + threadIdx.x;
+            if(i < n){
+                T val = output[i] + shortcut[i];
+                output[i] = val > T(0) ? val : T(0);
+            }
+        }
+        template<typename DEV_SPEC, typename T>
+        __global__
+        void d_relu_and_add(devices::CUDA<DEV_SPEC> device, const T* output, const T* d_output, T* d_pre_relu, unsigned int n){
+            unsigned int i = blockIdx.x * blockDim.x + threadIdx.x;
+            if(i < n){
+                d_pre_relu[i] = output[i] > T(0) ? d_output[i] : T(0);
+            }
+        }
+        template<typename DEV_SPEC, typename T>
+        __global__
+        void add_inplace(devices::CUDA<DEV_SPEC> device, T* dst, const T* src, unsigned int n){
+            unsigned int i = blockIdx.x * blockDim.x + threadIdx.x;
+            if(i < n){
+                dst[i] += src[i];
+            }
+        }
+    }
+    template<typename DEV_SPEC, typename LAYER_SPEC, typename INPUT_SPEC, typename OUTPUT_SPEC, typename BUFFER_SPEC, typename RNG, typename MODE = mode::Default<>>
+    void evaluate(devices::CUDA<DEV_SPEC>& device, const nn::layers::resnet_block::LayerForward<LAYER_SPEC>& layer, const Tensor<INPUT_SPEC>& input, Tensor<OUTPUT_SPEC>& output, nn::layers::resnet_block::Buffer<BUFFER_SPEC>& buffer, RNG& rng, const Mode<MODE>& mode = Mode<mode::Default<>>{}){
+        static_assert(nn::layers::resnet_block::check_input_output<LAYER_SPEC, INPUT_SPEC, OUTPUT_SPEC>);
+        using DEVICE = devices::CUDA<DEV_SPEC>;
+        using T = typename OUTPUT_SPEC::T;
+        using TI = typename DEVICE::index_t;
+        constexpr TI N = LAYER_SPEC::INTERNAL_BATCH_SIZE, OH = LAYER_SPEC::OUTPUT_HEIGHT, OW = LAYER_SPEC::OUTPUT_WIDTH, OC = LAYER_SPEC::OUTPUT_CHANNELS;
+        constexpr TI OUTPUT_SIZE = N * OH * OW * OC;
+        evaluate(device, static_cast<const nn::layers::conv2d::LayerForward<typename LAYER_SPEC::CONV1_SPEC>&>(layer.conv1), input, buffer.intermediate, buffer.conv1_buffer, rng, mode);
+        evaluate(device, static_cast<const nn::layers::conv2d::LayerForward<typename LAYER_SPEC::CONV2_SPEC>&>(layer.conv2), buffer.intermediate, output, buffer.conv2_buffer, rng, mode);
+        if constexpr(LAYER_SPEC::HAS_DOWNSAMPLE){
+            evaluate(device, static_cast<const nn::layers::conv2d::LayerForward<typename LAYER_SPEC::DOWNSAMPLE_SPEC>&>(layer.downsample.conv), input, buffer.shortcut, buffer.downsample_buffer, rng, mode);
+        }
+        constexpr TI BLOCKSIZE = 256;
+        constexpr TI N_BLOCKS = RL_TOOLS_DEVICES_CUDA_CEIL(OUTPUT_SIZE, BLOCKSIZE);
+        devices::cuda::TAG<DEVICE, true> tag_device{};
+        const T* shortcut_ptr = LAYER_SPEC::HAS_DOWNSAMPLE ? buffer.shortcut._data : input._data;
+        nn::layers::resnet_block::cuda::kernels::add_relu<<<N_BLOCKS, BLOCKSIZE, 0, device.stream>>>(tag_device, output._data, shortcut_ptr, OUTPUT_SIZE);
+        check_status(device);
+    }
+    template<typename DEV_SPEC, typename LAYER_SPEC, typename INPUT_SPEC, typename OUTPUT_SPEC, typename BUFFER_SPEC, typename RNG, typename MODE = mode::Default<>>
+    void forward(devices::CUDA<DEV_SPEC>& device, nn::layers::resnet_block::LayerBackward<LAYER_SPEC>& layer, const Tensor<INPUT_SPEC>& input, Tensor<OUTPUT_SPEC>& output, nn::layers::resnet_block::Buffer<BUFFER_SPEC>& buffer, RNG& rng, const Mode<MODE>& mode = Mode<mode::Default<>>{}){
+        static_assert(nn::layers::resnet_block::check_input_output<LAYER_SPEC, INPUT_SPEC, OUTPUT_SPEC>);
+        using DEVICE = devices::CUDA<DEV_SPEC>;
+        using T = typename OUTPUT_SPEC::T;
+        using TI = typename DEVICE::index_t;
+        constexpr TI N = LAYER_SPEC::INTERNAL_BATCH_SIZE, OH = LAYER_SPEC::OUTPUT_HEIGHT, OW = LAYER_SPEC::OUTPUT_WIDTH, OC = LAYER_SPEC::OUTPUT_CHANNELS;
+        constexpr TI OUTPUT_SIZE = N * OH * OW * OC;
+        forward(device, layer.conv1, input, buffer.intermediate, buffer.conv1_buffer, rng, mode);
+        forward(device, layer.conv2, buffer.intermediate, output, buffer.conv2_buffer, rng, mode);
+        if constexpr(LAYER_SPEC::HAS_DOWNSAMPLE){
+            forward(device, layer.downsample.conv, input, buffer.shortcut, buffer.downsample_buffer, rng, mode);
+        }
+        constexpr TI BLOCKSIZE = 256;
+        constexpr TI N_BLOCKS = RL_TOOLS_DEVICES_CUDA_CEIL(OUTPUT_SIZE, BLOCKSIZE);
+        devices::cuda::TAG<DEVICE, true> tag_device{};
+        const T* shortcut_ptr = LAYER_SPEC::HAS_DOWNSAMPLE ? buffer.shortcut._data : input._data;
+        nn::layers::resnet_block::cuda::kernels::add_relu<<<N_BLOCKS, BLOCKSIZE, 0, device.stream>>>(tag_device, output._data, shortcut_ptr, OUTPUT_SIZE);
+        check_status(device);
+    }
+    template<typename DEV_SPEC, typename LAYER_SPEC, typename INPUT_SPEC, typename D_OUTPUT_SPEC, typename BUFFER_SPEC, typename MODE = mode::Default<>>
+    void backward(devices::CUDA<DEV_SPEC>& device, nn::layers::resnet_block::LayerGradient<LAYER_SPEC>& layer, const Tensor<INPUT_SPEC>& input, Tensor<D_OUTPUT_SPEC>& d_output, nn::layers::resnet_block::Buffer<BUFFER_SPEC>& buffer, const Mode<MODE>& mode = Mode<mode::Default<>>{}){
+        using DEVICE = devices::CUDA<DEV_SPEC>;
+        using T = typename LAYER_SPEC::TYPE_POLICY::template GET<numeric_types::categories::Gradient>;
+        using TI = typename DEVICE::index_t;
+        constexpr TI N = LAYER_SPEC::INTERNAL_BATCH_SIZE, OH = LAYER_SPEC::OUTPUT_HEIGHT, OW = LAYER_SPEC::OUTPUT_WIDTH, OC = LAYER_SPEC::OUTPUT_CHANNELS;
+        constexpr TI OUTPUT_SIZE = N * OH * OW * OC;
+        constexpr TI BLOCKSIZE = 256;
+        constexpr TI N_BLOCKS = RL_TOOLS_DEVICES_CUDA_CEIL(OUTPUT_SIZE, BLOCKSIZE);
+        devices::cuda::TAG<DEVICE, true> tag_device{};
+        nn::layers::resnet_block::cuda::kernels::d_relu_and_add<<<N_BLOCKS, BLOCKSIZE, 0, device.stream>>>(tag_device, layer.output._data, d_output._data, buffer.shortcut._data, OUTPUT_SIZE);
+        auto conv1_out = rl_tools::output(device, layer.conv1);
+        backward(device, layer.conv2, conv1_out, buffer.shortcut, buffer.conv2_buffer, mode);
+        auto d_conv1_out = buffer.intermediate;
+        backward_input(device, static_cast<const nn::layers::conv2d::LayerBackward<typename LAYER_SPEC::CONV2_SPEC>&>(layer.conv2), buffer.shortcut, d_conv1_out, buffer.conv2_buffer, mode);
+        backward(device, layer.conv1, input, d_conv1_out, buffer.conv1_buffer, mode);
+        if constexpr(LAYER_SPEC::HAS_DOWNSAMPLE){
+            backward(device, layer.downsample.conv, input, buffer.shortcut, buffer.downsample_buffer, mode);
+        }
+    }
+    template<typename DEV_SPEC, typename LAYER_SPEC, typename INPUT_SPEC, typename D_OUTPUT_SPEC, typename D_INPUT_SPEC, typename BUFFER_SPEC, typename MODE = mode::Default<>>
+    void backward_full(devices::CUDA<DEV_SPEC>& device, nn::layers::resnet_block::LayerGradient<LAYER_SPEC>& layer, const Tensor<INPUT_SPEC>& input, Tensor<D_OUTPUT_SPEC>& d_output, Tensor<D_INPUT_SPEC>& d_input, nn::layers::resnet_block::Buffer<BUFFER_SPEC>& buffer, const Mode<MODE>& mode = Mode<mode::Default<>>{}){
+        using DEVICE = devices::CUDA<DEV_SPEC>;
+        using T = typename LAYER_SPEC::TYPE_POLICY::template GET<numeric_types::categories::Gradient>;
+        using TI = typename DEVICE::index_t;
+        constexpr TI N = LAYER_SPEC::INTERNAL_BATCH_SIZE;
+        constexpr TI OH = LAYER_SPEC::OUTPUT_HEIGHT, OW = LAYER_SPEC::OUTPUT_WIDTH, OC = LAYER_SPEC::OUTPUT_CHANNELS;
+        constexpr TI IH = LAYER_SPEC::INPUT_HEIGHT, IW = LAYER_SPEC::INPUT_WIDTH, IC = LAYER_SPEC::INPUT_CHANNELS;
+        constexpr TI OUTPUT_SIZE = N * OH * OW * OC;
+        constexpr TI INPUT_SIZE = N * IH * IW * IC;
+        constexpr TI BLOCKSIZE = 256;
+        devices::cuda::TAG<DEVICE, true> tag_device{};
+        // d_pre_relu = d_output * relu'(output)
+        nn::layers::resnet_block::cuda::kernels::d_relu_and_add<<<RL_TOOLS_DEVICES_CUDA_CEIL(OUTPUT_SIZE, BLOCKSIZE), BLOCKSIZE, 0, device.stream>>>(tag_device, layer.output._data, d_output._data, buffer.shortcut._data, OUTPUT_SIZE);
+        // backward_full through conv2
+        auto conv1_out = rl_tools::output(device, layer.conv1);
+        auto& d_pre_relu = buffer.shortcut;
+        backward_full(device, layer.conv2, conv1_out, d_pre_relu, buffer.intermediate, buffer.conv2_buffer, mode);
+        // backward_full through conv1
+        backward_full(device, layer.conv1, input, buffer.intermediate, d_input, buffer.conv1_buffer, mode);
+        // shortcut path
+        if constexpr(LAYER_SPEC::HAS_DOWNSAMPLE){
+            backward_full(device, layer.downsample.conv, input, d_pre_relu, buffer.d_input_buffer, buffer.downsample_buffer, mode);
+            nn::layers::resnet_block::cuda::kernels::add_inplace<<<RL_TOOLS_DEVICES_CUDA_CEIL(INPUT_SIZE, BLOCKSIZE), BLOCKSIZE, 0, device.stream>>>(tag_device, d_input._data, buffer.d_input_buffer._data, INPUT_SIZE);
+        } else {
+            nn::layers::resnet_block::cuda::kernels::add_inplace<<<RL_TOOLS_DEVICES_CUDA_CEIL(INPUT_SIZE, BLOCKSIZE), BLOCKSIZE, 0, device.stream>>>(tag_device, d_input._data, d_pre_relu._data, INPUT_SIZE);
+        }
+        check_status(device);
+    }
+    template<typename DEV_SPEC, typename SPEC>
+    void zero_gradient(devices::CUDA<DEV_SPEC>& device, nn::layers::resnet_block::LayerGradient<SPEC>& layer) {
+        zero_gradient(device, layer.conv1);
+        zero_gradient(device, layer.conv2);
+        zero_gradient(device, layer.downsample);
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+#endif

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

@@ -0,0 +1,428 @@
+#include "../../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_NN_LAYERS_RESNET_BLOCK_OPERATIONS_GENERIC_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_NN_LAYERS_RESNET_BLOCK_OPERATIONS_GENERIC_H
+#include "../../../containers/tensor/tensor.h"
+#include "../conv2d/operations_generic.h"
+#include "layer.h"
+#ifndef RL_TOOLS_FUNCTION_PLACEMENT
+#define RL_TOOLS_FUNCTION_PLACEMENT
+#endif
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools{
+    // ======================== DownsampleStorage malloc / free / copy / etc ========================
+    template<typename DEVICE, typename LAYER>
+    RL_TOOLS_FUNCTION_PLACEMENT void malloc(DEVICE&, nn::layers::resnet_block::DownsampleStorage<false, LAYER>&) {}
+    template<typename DEVICE, typename LAYER>
+    RL_TOOLS_FUNCTION_PLACEMENT void free(DEVICE&, nn::layers::resnet_block::DownsampleStorage<false, LAYER>&) {}
+    template<typename SOURCE_DEVICE, typename TARGET_DEVICE, typename SOURCE_LAYER, typename TARGET_LAYER>
+    RL_TOOLS_FUNCTION_PLACEMENT void copy(SOURCE_DEVICE&, TARGET_DEVICE&, const nn::layers::resnet_block::DownsampleStorage<false, SOURCE_LAYER>&, nn::layers::resnet_block::DownsampleStorage<false, TARGET_LAYER>&) {}
+    template<typename DEVICE, typename LAYER>
+    RL_TOOLS_FUNCTION_PLACEMENT void zero_gradient(DEVICE&, nn::layers::resnet_block::DownsampleStorage<false, LAYER>&) {}
+    template<typename DEVICE, typename LAYER, typename OPTIMIZER>
+    RL_TOOLS_FUNCTION_PLACEMENT void update(DEVICE&, nn::layers::resnet_block::DownsampleStorage<false, LAYER>&, OPTIMIZER&) {}
+    template<typename DEVICE, typename LAYER, typename OPTIMIZER>
+    RL_TOOLS_FUNCTION_PLACEMENT void _reset_optimizer_state(DEVICE&, nn::layers::resnet_block::DownsampleStorage<false, LAYER>&, OPTIMIZER&) {}
+    template<typename DEVICE, typename SOURCE_LAYER, typename TARGET_LAYER>
+    RL_TOOLS_FUNCTION_PLACEMENT typename SOURCE_LAYER::TYPE_POLICY::DEFAULT abs_diff(DEVICE&, const nn::layers::resnet_block::DownsampleStorage<false, SOURCE_LAYER>&, const nn::layers::resnet_block::DownsampleStorage<false, TARGET_LAYER>&) { return 0; }
+    template<typename DEVICE, typename LAYER, typename MODE>
+    RL_TOOLS_FUNCTION_PLACEMENT bool is_nan(DEVICE&, const nn::layers::resnet_block::DownsampleStorage<false, LAYER>&, const Mode<MODE>&) { return false; }
+    template<typename DEVICE, typename LAYER>
+    RL_TOOLS_FUNCTION_PLACEMENT void reset_forward_state(DEVICE&, nn::layers::resnet_block::DownsampleStorage<false, LAYER>&) {}
+    template<typename DEVICE, typename LAYER, typename RNG>
+    RL_TOOLS_FUNCTION_PLACEMENT void init_weights(DEVICE&, nn::layers::resnet_block::DownsampleStorage<false, LAYER>&, RNG&) {}
+
+    template<typename DEVICE, typename LAYER>
+    RL_TOOLS_FUNCTION_PLACEMENT void malloc(DEVICE& device, nn::layers::resnet_block::DownsampleStorage<true, LAYER>& ds) { malloc(device, ds.conv); }
+    template<typename DEVICE, typename LAYER>
+    RL_TOOLS_FUNCTION_PLACEMENT void free(DEVICE& device, nn::layers::resnet_block::DownsampleStorage<true, LAYER>& ds) { free(device, ds.conv); }
+    template<typename SOURCE_DEVICE, typename TARGET_DEVICE, typename SOURCE_LAYER, typename TARGET_LAYER>
+    RL_TOOLS_FUNCTION_PLACEMENT void copy(SOURCE_DEVICE& source_device, TARGET_DEVICE& target_device, const nn::layers::resnet_block::DownsampleStorage<true, SOURCE_LAYER>& source, nn::layers::resnet_block::DownsampleStorage<true, TARGET_LAYER>& target) { copy(source_device, target_device, source.conv, target.conv); }
+    template<typename DEVICE, typename LAYER>
+    RL_TOOLS_FUNCTION_PLACEMENT void zero_gradient(DEVICE& device, nn::layers::resnet_block::DownsampleStorage<true, LAYER>& ds) { zero_gradient(device, ds.conv); }
+    template<typename DEVICE, typename LAYER, typename OPTIMIZER>
+    RL_TOOLS_FUNCTION_PLACEMENT void update(DEVICE& device, nn::layers::resnet_block::DownsampleStorage<true, LAYER>& ds, OPTIMIZER& opt) { update(device, ds.conv, opt); }
+    template<typename DEVICE, typename LAYER, typename OPTIMIZER>
+    RL_TOOLS_FUNCTION_PLACEMENT void _reset_optimizer_state(DEVICE& device, nn::layers::resnet_block::DownsampleStorage<true, LAYER>& ds, OPTIMIZER& opt) { _reset_optimizer_state(device, ds.conv, opt); }
+    template<typename DEVICE, typename SOURCE_LAYER, typename TARGET_LAYER>
+    RL_TOOLS_FUNCTION_PLACEMENT auto abs_diff(DEVICE& device, const nn::layers::resnet_block::DownsampleStorage<true, SOURCE_LAYER>& a, const nn::layers::resnet_block::DownsampleStorage<true, TARGET_LAYER>& b) { return abs_diff(device, a.conv, b.conv); }
+    template<typename DEVICE, typename LAYER, typename MODE>
+    RL_TOOLS_FUNCTION_PLACEMENT bool is_nan(DEVICE& device, const nn::layers::resnet_block::DownsampleStorage<true, LAYER>& ds, const Mode<MODE>& mode) { return is_nan(device, ds.conv, mode); }
+    template<typename DEVICE, typename LAYER>
+    RL_TOOLS_FUNCTION_PLACEMENT void reset_forward_state(DEVICE& device, nn::layers::resnet_block::DownsampleStorage<true, LAYER>& ds) { reset_forward_state(device, ds.conv); }
+    template<typename DEVICE, typename LAYER, typename RNG>
+    RL_TOOLS_FUNCTION_PLACEMENT void init_weights(DEVICE& device, nn::layers::resnet_block::DownsampleStorage<true, LAYER>& ds, RNG& rng) { init_weights(device, ds.conv, rng); }
+
+    // ======================== Buffer malloc / free ========================
+    template<typename DEVICE, typename BUFFER_SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void malloc(DEVICE& device, nn::layers::resnet_block::Buffer<BUFFER_SPEC>& buffer) {
+        malloc(device, buffer.intermediate);
+        malloc(device, buffer.shortcut);
+        malloc(device, buffer.d_input_buffer);
+        malloc(device, buffer.conv1_buffer);
+        malloc(device, buffer.conv2_buffer);
+        malloc(device, buffer.downsample_buffer);
+    }
+    template<typename DEVICE, typename BUFFER_SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void free(DEVICE& device, nn::layers::resnet_block::Buffer<BUFFER_SPEC>& buffer) {
+        free(device, buffer.intermediate);
+        free(device, buffer.shortcut);
+        free(device, buffer.d_input_buffer);
+        free(device, buffer.conv1_buffer);
+        free(device, buffer.conv2_buffer);
+        free(device, buffer.downsample_buffer);
+    }
+
+    // ======================== Layer malloc / free ========================
+    template<typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void malloc(DEVICE& device, nn::layers::resnet_block::LayerForward<SPEC>& layer) {
+        malloc(device, layer.conv1);
+        malloc(device, layer.conv2);
+        malloc(device, layer.downsample);
+    }
+    template<typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void free(DEVICE& device, nn::layers::resnet_block::LayerForward<SPEC>& layer) {
+        free(device, layer.conv1);
+        free(device, layer.conv2);
+        free(device, layer.downsample);
+    }
+    template<typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void malloc(DEVICE& device, nn::layers::resnet_block::LayerBackward<SPEC>& layer) {
+        malloc(device, (nn::layers::resnet_block::LayerForward<SPEC>&)layer);
+    }
+    template<typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void free(DEVICE& device, nn::layers::resnet_block::LayerBackward<SPEC>& layer) {
+        free(device, (nn::layers::resnet_block::LayerForward<SPEC>&)layer);
+    }
+    template<typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void malloc(DEVICE& device, nn::layers::resnet_block::LayerGradient<SPEC>& layer) {
+        malloc(device, (nn::layers::resnet_block::LayerBackward<SPEC>&)layer);
+        malloc(device, layer.output);
+    }
+    template<typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void free(DEVICE& device, nn::layers::resnet_block::LayerGradient<SPEC>& layer) {
+        free(device, (nn::layers::resnet_block::LayerBackward<SPEC>&)layer);
+        free(device, layer.output);
+    }
+    template<typename DEVICE>
+    RL_TOOLS_FUNCTION_PLACEMENT void malloc(DEVICE&, nn::layers::resnet_block::State&) {}
+    template<typename DEVICE>
+    RL_TOOLS_FUNCTION_PLACEMENT void free(DEVICE&, nn::layers::resnet_block::State&) {}
+    template<typename SOURCE_DEVICE, typename TARGET_DEVICE>
+    RL_TOOLS_FUNCTION_PLACEMENT void copy(SOURCE_DEVICE&, TARGET_DEVICE&, nn::layers::resnet_block::State&, nn::layers::resnet_block::State&) {}
+    template<typename SPEC, typename DEVICE, typename RNG, typename MODE>
+    RL_TOOLS_FUNCTION_PLACEMENT void reset(DEVICE&, const nn::layers::resnet_block::LayerForward<SPEC>&, nn::layers::resnet_block::State&, RNG&, Mode<MODE>) {}
+
+    // ======================== init_weights ========================
+    template<typename DEVICE, typename SPEC, typename RNG>
+    RL_TOOLS_FUNCTION_PLACEMENT void init_weights(DEVICE& device, nn::layers::resnet_block::LayerForward<SPEC>& layer, RNG& rng) {
+        init_weights(device, layer.conv1, rng);
+        init_weights(device, layer.conv2, rng);
+        init_weights(device, layer.downsample, rng);
+    }
+
+    // ======================== evaluate ========================
+#ifndef RL_TOOLS_NN_DISABLE_GENERIC_FORWARD_BACKWARD
+    template<typename DEVICE, typename LAYER_SPEC, typename INPUT_SPEC, typename OUTPUT_SPEC, typename BUFFER_SPEC, typename RNG, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT void evaluate(DEVICE& device, const nn::layers::resnet_block::LayerForward<LAYER_SPEC>& layer, const Tensor<INPUT_SPEC>& input, Tensor<OUTPUT_SPEC>& output, nn::layers::resnet_block::Buffer<BUFFER_SPEC>& buffer, RNG& rng, const Mode<MODE>& mode = Mode<mode::Default<>>{}) {
+        static_assert(nn::layers::resnet_block::check_input_output<LAYER_SPEC, INPUT_SPEC, OUTPUT_SPEC>);
+        using TI = typename DEVICE::index_t;
+        using T = typename OUTPUT_SPEC::T;
+        constexpr TI BATCH = LAYER_SPEC::INTERNAL_BATCH_SIZE;
+        constexpr TI OH = LAYER_SPEC::OUTPUT_HEIGHT;
+        constexpr TI OW = LAYER_SPEC::OUTPUT_WIDTH;
+        constexpr TI OC = LAYER_SPEC::OUTPUT_CHANNELS;
+
+        using INTERNAL_INPUT_SHAPE = tensor::Shape<TI, BATCH, LAYER_SPEC::INPUT_HEIGHT, LAYER_SPEC::INPUT_WIDTH, LAYER_SPEC::INPUT_CHANNELS>;
+        using INTERNAL_OUTPUT_SHAPE = tensor::Shape<TI, BATCH, OH, OW, OC>;
+        auto input_4d = view_memory<INTERNAL_INPUT_SHAPE>(device, input);
+        auto output_4d = view_memory<INTERNAL_OUTPUT_SHAPE>(device, output);
+
+        // Conv1: input → intermediate
+        evaluate(device, static_cast<const nn::layers::conv2d::LayerForward<typename LAYER_SPEC::CONV1_SPEC>&>(layer.conv1), input, buffer.intermediate, buffer.conv1_buffer, rng, mode);
+        // Conv2: intermediate → output (temporary, before skip)
+        evaluate(device, static_cast<const nn::layers::conv2d::LayerForward<typename LAYER_SPEC::CONV2_SPEC>&>(layer.conv2), buffer.intermediate, output, buffer.conv2_buffer, rng, mode);
+        // Downsample shortcut if needed
+        if constexpr(LAYER_SPEC::HAS_DOWNSAMPLE) {
+            evaluate(device, static_cast<const nn::layers::conv2d::LayerForward<typename LAYER_SPEC::DOWNSAMPLE_SPEC>&>(layer.downsample.conv), input, buffer.shortcut, buffer.downsample_buffer, rng, mode);
+        }
+        // output = ReLU(conv2_out + shortcut)
+        for(TI bi = 0; bi < BATCH; bi++){
+            for(TI h = 0; h < OH; h++){
+                for(TI w = 0; w < OW; w++){
+                    for(TI c = 0; c < OC; c++){
+                        T conv2_val = get(device, output_4d, bi, h, w, c);
+                        T shortcut_val;
+                        if constexpr(LAYER_SPEC::HAS_DOWNSAMPLE) {
+                            shortcut_val = get(device, buffer.shortcut, bi, h, w, c);
+                        } else {
+                            shortcut_val = get(device, input_4d, bi, h, w, c);
+                        }
+                        set(device, output_4d, math::max(device.math, conv2_val + shortcut_val, (T)0), bi, h, w, c);
+                    }
+                }
+            }
+        }
+    }
+
+    template<typename DEVICE, typename LAYER_SPEC, typename INPUT_SPEC, typename OUTPUT_SPEC, typename BUFFER_SPEC, typename RNG, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT void forward(DEVICE& device, nn::layers::resnet_block::LayerBackward<LAYER_SPEC>& layer, const Tensor<INPUT_SPEC>& input, Tensor<OUTPUT_SPEC>& output, nn::layers::resnet_block::Buffer<BUFFER_SPEC>& buffer, RNG& rng, const Mode<MODE>& mode = Mode<mode::Default<>>{}){
+        static_assert(nn::layers::resnet_block::check_input_output<LAYER_SPEC, INPUT_SPEC, OUTPUT_SPEC>);
+        using TI = typename DEVICE::index_t;
+        using T = typename OUTPUT_SPEC::T;
+        constexpr TI BATCH = LAYER_SPEC::INTERNAL_BATCH_SIZE;
+        constexpr TI OH = LAYER_SPEC::OUTPUT_HEIGHT;
+        constexpr TI OW = LAYER_SPEC::OUTPUT_WIDTH;
+        constexpr TI OC = LAYER_SPEC::OUTPUT_CHANNELS;
+
+        using INTERNAL_INPUT_SHAPE = tensor::Shape<TI, BATCH, LAYER_SPEC::INPUT_HEIGHT, LAYER_SPEC::INPUT_WIDTH, LAYER_SPEC::INPUT_CHANNELS>;
+        auto input_4d = view_memory<INTERNAL_INPUT_SHAPE>(device, input);
+        auto output_4d = view_memory<tensor::Shape<TI, BATCH, OH, OW, OC>>(device, output);
+
+        // Conv1: input → buffer.intermediate
+        forward(device, layer.conv1, input, buffer.intermediate, buffer.conv1_buffer, rng, mode);
+        // Conv2: buffer.intermediate → output (before skip addition)
+        forward(device, layer.conv2, buffer.intermediate, output, buffer.conv2_buffer, rng, mode);
+        // Downsample shortcut if needed
+        if constexpr(LAYER_SPEC::HAS_DOWNSAMPLE) {
+            forward(device, layer.downsample.conv, input, buffer.shortcut, buffer.downsample_buffer, rng, mode);
+        }
+        // output = ReLU(conv2_out + shortcut)
+        for(TI bi = 0; bi < BATCH; bi++){
+            for(TI h = 0; h < OH; h++){
+                for(TI w = 0; w < OW; w++){
+                    for(TI c = 0; c < OC; c++){
+                        T conv2_val = get(device, output_4d, bi, h, w, c);
+                        T shortcut_val;
+                        if constexpr(LAYER_SPEC::HAS_DOWNSAMPLE) {
+                            shortcut_val = get(device, buffer.shortcut, bi, h, w, c);
+                        } else {
+                            shortcut_val = get(device, input_4d, bi, h, w, c);
+                        }
+                        set(device, output_4d, math::max(device.math, conv2_val + shortcut_val, (T)0), bi, h, w, c);
+                    }
+                }
+            }
+        }
+    }
+#endif
+
+    template<typename DEVICE, typename LAYER_SPEC, typename INPUT_SPEC, typename BUFFER_SPEC, typename RNG, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT void forward(DEVICE& device, nn::layers::resnet_block::LayerGradient<LAYER_SPEC>& layer, const Tensor<INPUT_SPEC>& input, nn::layers::resnet_block::Buffer<BUFFER_SPEC>& buffer, RNG& rng, const Mode<MODE>& mode = Mode<mode::Default<>>{}) {
+        forward(device, static_cast<nn::layers::resnet_block::LayerBackward<LAYER_SPEC>&>(layer), input, layer.output, buffer, rng, mode);
+    }
+    template<typename DEVICE, typename LAYER_SPEC, typename INPUT_SPEC, typename OUTPUT_SPEC, typename BUFFER_SPEC, typename RNG, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT void forward(DEVICE& device, nn::layers::resnet_block::LayerGradient<LAYER_SPEC>& layer, const Tensor<INPUT_SPEC>& input, Tensor<OUTPUT_SPEC>& output, nn::layers::resnet_block::Buffer<BUFFER_SPEC>& buffer, RNG& rng, const Mode<MODE>& mode = Mode<mode::Default<>>{}) {
+        forward(device, layer, input, buffer, rng, mode);
+        copy(device, device, layer.output, output);
+    }
+
+    // ======================== backward ========================
+    // backward_input: compute d_input only (no gradient accumulation)
+    template<typename DEVICE, typename LAYER_SPEC, typename D_OUTPUT_SPEC, typename D_INPUT_SPEC, typename BUFFER_SPEC, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT void backward_input(DEVICE& device, const nn::layers::resnet_block::LayerBackward<LAYER_SPEC>& layer, const Tensor<D_OUTPUT_SPEC>& d_output, Tensor<D_INPUT_SPEC>& d_input, nn::layers::resnet_block::Buffer<BUFFER_SPEC>& buffer, const Mode<MODE>& mode = Mode<mode::Default<>>{}){
+        // Not implemented standalone (use backward_full)
+    }
+    // backward: accumulate parameter gradients only (no d_input)
+    template<typename DEVICE, typename LAYER_SPEC, typename INPUT_SPEC, typename D_OUTPUT_SPEC, typename BUFFER_SPEC, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT void backward(DEVICE& device, nn::layers::resnet_block::LayerGradient<LAYER_SPEC>& layer, const Tensor<INPUT_SPEC>& input, Tensor<D_OUTPUT_SPEC>& d_output, nn::layers::resnet_block::Buffer<BUFFER_SPEC>& buffer, const Mode<MODE>& mode = Mode<mode::Default<>>{}){
+        using TI = typename DEVICE::index_t;
+        using T = typename LAYER_SPEC::TYPE_POLICY::template GET<numeric_types::categories::Gradient>;
+        constexpr TI BATCH = LAYER_SPEC::INTERNAL_BATCH_SIZE;
+        constexpr TI OH = LAYER_SPEC::OUTPUT_HEIGHT;
+        constexpr TI OW = LAYER_SPEC::OUTPUT_WIDTH;
+        constexpr TI OC = LAYER_SPEC::OUTPUT_CHANNELS;
+        // Step 1: d_pre_relu = d_output * relu'(output)
+        auto output_view = view_memory<tensor::Shape<TI, BATCH, OH, OW, OC>>(device, layer.output);
+        auto d_output_4d = view_memory<tensor::Shape<TI, BATCH, OH, OW, OC>>(device, d_output);
+        auto d_pre_relu = view_memory<tensor::Shape<TI, BATCH, OH, OW, OC>>(device, buffer.shortcut);
+        for(TI bi = 0; bi < BATCH; bi++){
+            for(TI h = 0; h < OH; h++){
+                for(TI w = 0; w < OW; w++){
+                    for(TI c = 0; c < OC; c++){
+                        T out_val = get(device, output_view, bi, h, w, c);
+                        T d_val = get(device, d_output_4d, bi, h, w, c);
+                        set(device, d_pre_relu, out_val > (T)0 ? d_val : (T)0, bi, h, w, c);
+                    }
+                }
+            }
+        }
+        // Step 2: backward through conv2 (accumulate gradients, get d_conv1_out)
+        auto conv1_out = rl_tools::output(device, layer.conv1);
+        backward(device, layer.conv2, conv1_out, d_pre_relu, buffer.conv2_buffer, mode);
+        // Step 3: backward through conv1 (accumulate gradients)
+        // Need d_conv1_out from conv2's backward_input
+        auto d_conv1_out = view_memory<tensor::Shape<TI, BATCH, OH, OW, OC>>(device, buffer.intermediate);
+        backward_input(device, static_cast<const nn::layers::conv2d::LayerBackward<typename LAYER_SPEC::CONV2_SPEC>&>(layer.conv2), d_pre_relu, d_conv1_out, buffer.conv2_buffer, mode);
+        backward(device, layer.conv1, input, d_conv1_out, buffer.conv1_buffer, mode);
+        // Step 4: backward through downsample (accumulate gradients)
+        if constexpr(LAYER_SPEC::HAS_DOWNSAMPLE) {
+            backward(device, layer.downsample.conv, input, d_pre_relu, buffer.downsample_buffer, mode);
+        }
+    }
+    // backward_full: compute d_input AND accumulate parameter gradients
+    template<typename DEVICE, typename LAYER_SPEC, typename INPUT_SPEC, typename D_OUTPUT_SPEC, typename D_INPUT_SPEC, typename BUFFER_SPEC, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT void backward_full(DEVICE& device, nn::layers::resnet_block::LayerGradient<LAYER_SPEC>& layer, const Tensor<INPUT_SPEC>& input, Tensor<D_OUTPUT_SPEC>& d_output, Tensor<D_INPUT_SPEC>& d_input, nn::layers::resnet_block::Buffer<BUFFER_SPEC>& buffer, const Mode<MODE>& mode = Mode<mode::Default<>>{}){
+        using TI = typename DEVICE::index_t;
+        using T = typename LAYER_SPEC::TYPE_POLICY::template GET<numeric_types::categories::Gradient>;
+        constexpr TI BATCH = LAYER_SPEC::INTERNAL_BATCH_SIZE;
+        constexpr TI OH = LAYER_SPEC::OUTPUT_HEIGHT;
+        constexpr TI OW = LAYER_SPEC::OUTPUT_WIDTH;
+        constexpr TI OC = LAYER_SPEC::OUTPUT_CHANNELS;
+        constexpr TI IH = LAYER_SPEC::INPUT_HEIGHT;
+        constexpr TI IW = LAYER_SPEC::INPUT_WIDTH;
+        constexpr TI IC = LAYER_SPEC::INPUT_CHANNELS;
+        // Step 1: d_pre_relu = d_output * relu'(output)
+        auto output_view = view_memory<tensor::Shape<TI, BATCH, OH, OW, OC>>(device, layer.output);
+        auto d_output_4d = view_memory<tensor::Shape<TI, BATCH, OH, OW, OC>>(device, d_output);
+        auto d_pre_relu = view_memory<tensor::Shape<TI, BATCH, OH, OW, OC>>(device, buffer.shortcut);
+        for(TI bi = 0; bi < BATCH; bi++){
+            for(TI h = 0; h < OH; h++){
+                for(TI w = 0; w < OW; w++){
+                    for(TI c = 0; c < OC; c++){
+                        T out_val = get(device, output_view, bi, h, w, c);
+                        T d_val = get(device, d_output_4d, bi, h, w, c);
+                        set(device, d_pre_relu, out_val > (T)0 ? d_val : (T)0, bi, h, w, c);
+                    }
+                }
+            }
+        }
+        // Step 2: backward_full through conv2 (accumulate gradients + d_conv1_out)
+        auto conv1_out = rl_tools::output(device, layer.conv1);
+        auto d_conv1_out = view_memory<tensor::Shape<TI, BATCH, OH, OW, OC>>(device, buffer.intermediate);
+        backward_full(device, layer.conv2, conv1_out, d_pre_relu, d_conv1_out, buffer.conv2_buffer, mode);
+        // Step 3: backward_full through conv1 (accumulate gradients + d_input)
+        backward_full(device, layer.conv1, input, d_conv1_out, d_input, buffer.conv1_buffer, mode);
+        // Step 4: backward through shortcut path and add to d_input
+        if constexpr(LAYER_SPEC::HAS_DOWNSAMPLE) {
+            auto d_input_ds = view_memory<tensor::Shape<TI, BATCH, IH, IW, IC>>(device, buffer.d_input_buffer);
+            backward_full(device, layer.downsample.conv, input, d_pre_relu, d_input_ds, buffer.downsample_buffer, mode);
+            // d_input += d_input_ds
+            auto d_input_4d = view_memory<tensor::Shape<TI, BATCH, IH, IW, IC>>(device, d_input);
+            for(TI bi = 0; bi < BATCH; bi++){
+                for(TI h = 0; h < IH; h++){
+                    for(TI w = 0; w < IW; w++){
+                        for(TI c = 0; c < IC; c++){
+                            increment(device, d_input_4d, get(device, d_input_ds, bi, h, w, c), bi, h, w, c);
+                        }
+                    }
+                }
+            }
+        } else {
+            // Identity shortcut: d_input += d_pre_relu
+            // d_pre_relu has shape (BATCH, OH, OW, OC) and d_input has shape (BATCH, IH, IW, IC)
+            // For identity shortcut, OH==IH, OW==IW, OC==IC
+            auto d_input_4d = view_memory<tensor::Shape<TI, BATCH, IH, IW, IC>>(device, d_input);
+            for(TI bi = 0; bi < BATCH; bi++){
+                for(TI h = 0; h < IH; h++){
+                    for(TI w = 0; w < IW; w++){
+                        for(TI c = 0; c < IC; c++){
+                            increment(device, d_input_4d, get(device, d_pre_relu, bi, h, w, c), bi, h, w, c);
+                        }
+                    }
+                }
+            }
+        }
+    }
+
+    // ======================== zero_gradient / update / _reset_optimizer_state ========================
+    template<typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void zero_gradient(DEVICE& device, nn::layers::resnet_block::LayerGradient<SPEC>& layer) {
+        zero_gradient(device, layer.conv1);
+        zero_gradient(device, layer.conv2);
+        zero_gradient(device, layer.downsample);
+    }
+    template<typename DEVICE, typename SPEC, typename OPTIMIZER>
+    RL_TOOLS_FUNCTION_PLACEMENT void update(DEVICE& device, nn::layers::resnet_block::LayerGradient<SPEC>& layer, OPTIMIZER& optimizer) {
+        update(device, layer.conv1, optimizer);
+        update(device, layer.conv2, optimizer);
+        update(device, layer.downsample, optimizer);
+    }
+    template<typename DEVICE, typename SPEC, typename OPTIMIZER>
+    RL_TOOLS_FUNCTION_PLACEMENT void _reset_optimizer_state(DEVICE& device, nn::layers::resnet_block::LayerGradient<SPEC>& layer, OPTIMIZER& optimizer) {
+        _reset_optimizer_state(device, layer.conv1, optimizer);
+        _reset_optimizer_state(device, layer.conv2, optimizer);
+        _reset_optimizer_state(device, layer.downsample, 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::resnet_block::LayerForward<SOURCE_SPEC>& source, nn::layers::resnet_block::LayerForward<TARGET_SPEC>& target) {
+        static_assert(nn::layers::resnet_block::check_spec_memory<SOURCE_SPEC, TARGET_SPEC>);
+        copy(source_device, target_device, source.conv1, target.conv1);
+        copy(source_device, target_device, source.conv2, target.conv2);
+        copy(source_device, target_device, source.downsample, target.downsample);
+    }
+    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::resnet_block::LayerBackward<SOURCE_SPEC>& source, nn::layers::resnet_block::LayerBackward<TARGET_SPEC>& target) {
+        static_assert(nn::layers::resnet_block::check_spec_memory<SOURCE_SPEC, TARGET_SPEC>);
+        copy(source_device, target_device, static_cast<const nn::layers::resnet_block::LayerForward<SOURCE_SPEC>&>(source), static_cast<nn::layers::resnet_block::LayerForward<TARGET_SPEC>&>(target));
+    }
+    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::resnet_block::LayerGradient<SOURCE_SPEC>& source, nn::layers::resnet_block::LayerGradient<TARGET_SPEC>& target) {
+        static_assert(nn::layers::resnet_block::check_spec_memory<SOURCE_SPEC, TARGET_SPEC>);
+        copy(source_device, target_device, static_cast<const nn::layers::resnet_block::LayerBackward<SOURCE_SPEC>&>(source), static_cast<nn::layers::resnet_block::LayerBackward<TARGET_SPEC>&>(target));
+        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::resnet_block::LayerForward<S1>& l1, const nn::layers::resnet_block::LayerForward<S2>& l2) {
+        using T = typename S1::TYPE_POLICY::DEFAULT;
+        T acc = 0;
+        acc += abs_diff(device, l1.conv1, l2.conv1);
+        acc += abs_diff(device, l1.conv2, l2.conv2);
+        acc += abs_diff(device, l1.downsample, l2.downsample);
+        return acc;
+    }
+    template<typename DEVICE, typename S1, typename S2>
+    RL_TOOLS_FUNCTION_PLACEMENT typename S1::TYPE_POLICY::DEFAULT abs_diff(DEVICE& device, const nn::layers::resnet_block::LayerBackward<S1>& l1, const nn::layers::resnet_block::LayerBackward<S2>& l2) {
+        return abs_diff(device, static_cast<const nn::layers::resnet_block::LayerForward<S1>&>(l1), static_cast<const nn::layers::resnet_block::LayerForward<S2>&>(l2));
+    }
+    template<typename DEVICE, typename S1, typename S2>
+    RL_TOOLS_FUNCTION_PLACEMENT typename S1::TYPE_POLICY::DEFAULT abs_diff(DEVICE& device, const nn::layers::resnet_block::LayerGradient<S1>& l1, const nn::layers::resnet_block::LayerGradient<S2>& l2) {
+        auto acc = abs_diff(device, static_cast<const nn::layers::resnet_block::LayerBackward<S1>&>(l1), static_cast<const nn::layers::resnet_block::LayerBackward<S2>&>(l2));
+        acc += abs_diff(device, l1.output, l2.output);
+        return acc;
+    }
+    template<typename DEVICE>
+    RL_TOOLS_FUNCTION_PLACEMENT auto abs_diff(DEVICE& device, const nn::layers::resnet_block::State&, const nn::layers::resnet_block::State&) { return 0; }
+
+    // ======================== reset_forward_state ========================
+    template<typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void reset_forward_state(DEVICE& device, nn::layers::resnet_block::LayerBackward<SPEC>& l) {
+        reset_forward_state(device, l.conv1);
+        reset_forward_state(device, l.conv2);
+        reset_forward_state(device, l.downsample);
+    }
+    template<typename DEVICE, typename SPEC>
+    RL_TOOLS_FUNCTION_PLACEMENT void reset_forward_state(DEVICE& device, nn::layers::resnet_block::LayerGradient<SPEC>& l) {
+        reset_forward_state(device, static_cast<nn::layers::resnet_block::LayerBackward<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& device, const nn::layers::resnet_block::LayerForward<SPEC>& l, const Mode<MODE>& mode = Mode<mode::Default<>>{}) {
+        return is_nan(device, l.conv1, mode) || is_nan(device, l.conv2, mode) || is_nan(device, l.downsample, mode);
+    }
+    template<typename DEVICE, typename SPEC, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT bool is_nan(DEVICE& device, const nn::layers::resnet_block::LayerBackward<SPEC>& l, const Mode<MODE>& mode = Mode<mode::Default<>>{}) {
+        return is_nan(device, static_cast<const nn::layers::resnet_block::LayerForward<SPEC>&>(l), mode);
+    }
+    template<typename DEVICE, typename SPEC, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT bool is_nan(DEVICE& device, const nn::layers::resnet_block::LayerGradient<SPEC>& l, const Mode<MODE>& mode = Mode<mode::Default<>>{}) {
+        return is_nan(device, static_cast<const nn::layers::resnet_block::LayerBackward<SPEC>&>(l), mode) || is_nan(device, l.output, mode);
+    }
+    template<typename DEVICE, typename MODE = mode::Default<>>
+    RL_TOOLS_FUNCTION_PLACEMENT bool is_nan(DEVICE&, nn::layers::resnet_block::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::resnet_block::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/layers/resnet_block/persist.h

@@ -0,0 +1,70 @@
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_NN_LAYERS_RESNET_BLOCK_PERSIST_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_NN_LAYERS_RESNET_BLOCK_PERSIST_H
+#include "../../../version.h"
+#include "layer.h"
+#include "../conv2d/persist.h"
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools {
+    // ======================== DownsampleStorage save / load ========================
+    template<typename DEVICE, typename LAYER, typename GROUP>
+    void save(DEVICE& device, nn::layers::resnet_block::DownsampleStorage<false, LAYER>&, GROUP& group) {}
+    template<typename DEVICE, typename LAYER, typename GROUP>
+    bool load(DEVICE& device, nn::layers::resnet_block::DownsampleStorage<false, LAYER>&, GROUP& group) { return true; }
+    template<typename DEVICE, typename LAYER, typename GROUP>
+    void save(DEVICE& device, nn::layers::resnet_block::DownsampleStorage<true, LAYER>& ds, GROUP& group) {
+        auto ds_group = create_group(device, group, "downsample");
+        save(device, ds.conv, ds_group);
+    }
+    template<typename DEVICE, typename LAYER, typename GROUP>
+    bool load(DEVICE& device, nn::layers::resnet_block::DownsampleStorage<true, LAYER>& ds, GROUP& group) {
+        if(group_exists(device, group, "downsample")){
+            auto ds_group = get_group(device, group, "downsample");
+            return load(device, ds.conv, ds_group);
+        }
+        return true;
+    }
+
+    // ======================== Layer save / load ========================
+    template<typename DEVICE, typename SPEC, typename GROUP>
+    void save(DEVICE& device, nn::layers::resnet_block::LayerForward<SPEC>& layer, GROUP& group) {
+        set_attribute(device, group, "type", "resnet_block");
+        write_attributes(device, group);
+        auto conv1_group = create_group(device, group, "conv1");
+        save(device, layer.conv1, conv1_group);
+        auto conv2_group = create_group(device, group, "conv2");
+        save(device, layer.conv2, conv2_group);
+        save(device, layer.downsample, group);
+    }
+    template<typename DEVICE, typename SPEC, typename GROUP>
+    void save(DEVICE& device, nn::layers::resnet_block::LayerBackward<SPEC>& layer, GROUP& group) {
+        save(device, (nn::layers::resnet_block::LayerForward<SPEC>&)layer, group);
+    }
+    template<typename DEVICE, typename SPEC, typename GROUP>
+    void save(DEVICE& device, nn::layers::resnet_block::LayerGradient<SPEC>& layer, GROUP& group) {
+        save(device, (nn::layers::resnet_block::LayerBackward<SPEC>&)layer, group);
+    }
+    template<typename DEVICE, typename SPEC, typename GROUP>
+    bool load(DEVICE& device, nn::layers::resnet_block::LayerForward<SPEC>& layer, GROUP& group) {
+        auto conv1_group = get_group(device, group, "conv1");
+        bool success = load(device, layer.conv1, conv1_group);
+        auto conv2_group = get_group(device, group, "conv2");
+        success &= load(device, layer.conv2, conv2_group);
+        success &= load(device, layer.downsample, group);
+        return success;
+    }
+    template<typename DEVICE, typename SPEC, typename GROUP>
+    bool load(DEVICE& device, nn::layers::resnet_block::LayerBackward<SPEC>& layer, GROUP& group) {
+        return load(device, (nn::layers::resnet_block::LayerForward<SPEC>&)layer, group);
+    }
+    template<typename DEVICE, typename SPEC, typename GROUP>
+    bool load(DEVICE& device, nn::layers::resnet_block::LayerGradient<SPEC>& layer, GROUP& group) {
+        return load(device, (nn::layers::resnet_block::LayerBackward<SPEC>&)layer, group);
+    }
+    template<typename DEVICE, typename GROUP>
+    void save(DEVICE& device, nn::layers::resnet_block::State& state, GROUP& group) {}
+    template<typename DEVICE, typename GROUP>
+    bool load(DEVICE& device, nn::layers::resnet_block::State& state, GROUP& group) { return true; }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+#endif