rustyml 0.17.0

//! 2D depthwise convolution layer
//!
//! [`DepthwiseConv2D`] gives each input channel its own kernel or kernels, with no mixing
//! between channels. `DepthwiseConv2DCache` holds the tensors the layer parks between its
//! forward and backward passes.

use crate::error::Error;
use crate::neural_network::layers::ParamCounts;
use crate::neural_network::layers::activation::{Activation, ActivationCache};
use crate::neural_network::layers::conv_op_helpers::{
    DepthwiseGeometry, depthwise_backward, depthwise_forward,
};
use crate::neural_network::layers::convolution::PaddingType;
use crate::neural_network::layers::convolution::convolution_engine::effective_kernel;
use crate::neural_network::layers::convolution::validation::{
    valid_output_size, validate_depth_multiplier, validate_dilation, validate_input_shape_2d,
    validate_kernel_size_2d, validate_strides_2d, validate_valid_kernel_fits,
};
use crate::neural_network::layers::shape_helpers::calculate_output_shape_2d;
use crate::neural_network::layers::validation::{
    start_build, validate_built_input, validate_optional_weight, validate_weight_shape,
};
use crate::neural_network::layers::{built_layer_shape_functions, named_weight_layer_functions};
use crate::neural_network::traits::{LayerBase, ParamRef, UnaryLayer};
use crate::neural_network::{Ctx, Fans, Initializer, Shape, Tensor};
use ndarray::{Array1, Array4};

/// A 2D depthwise convolutional layer
///
/// Applies its own kernel to each input channel. This lowers the parameter count and the
/// computation compared to a standard convolution. It still extracts spatial features per
/// channel. Input shape is \[batch_size, height, width, channels\] and output shape is
/// \[batch_size, out_height, out_width, channels * depth_multiplier\]. No kernel mixes channels
///
/// `depth_multiplier` (default 1, set with [`DepthwiseConv2D::with_depth_multiplier`]) is how many
/// kernels each input channel gets. The output channel for input channel `c` and multiplier index
/// `m` is `c * depth_multiplier + m`
///
/// # Examples
///
/// ```rust
/// use rustyml::neural_network::Shape;
/// use rustyml::neural_network::sequential::SequentialBuilder;
/// use rustyml::neural_network::layers::*;
/// use rustyml::neural_network::optimizers::*;
/// use rustyml::neural_network::losses::*;
/// use ndarray::Array4;
///
/// // Input data shape: [batch_size, height, width, channels].
/// let batch_size = 1;
/// let input_channels = 3;
/// let height = 4;
/// let width = 4;
///
/// let mut input_data = Array4::zeros((batch_size, height, width, input_channels));
///
/// // Give each channel distinct values.
/// for c in 0..input_channels {
///     for h in 0..height {
///         for w in 0..width {
///             input_data[[0, h, w, c]] = (c * 10 + h * width + w) as f32;
///         }
///     }
/// }
///
/// let input = input_data.into_dyn();
///
/// // DepthwiseConv2D with ReLU. The build draws the weights.
/// let depthwise_layer = DepthwiseConv2D::new(
///     (2, 2),                  // kernel_size
///     (1, 1),                  // strides
///     Activation::ReLU,        // activation
/// ).unwrap();
///
/// let mut model = SequentialBuilder::new()
///     .add(depthwise_layer)
///     .build(&Shape::known(input.shape()))
///     .unwrap();
/// model.compile(SGD::new(0.01, 0.0, false, 0.0).unwrap(), MeanSquaredError::new());
///
/// model.summary();
///
/// let output = model.predict(&input).unwrap();
///
/// // Input [1, 4, 4, 3] with kernel (2, 2), stride (1, 1), and Valid padding gives output.
/// // [1, 3, 3, 3].
/// assert_eq!(output.shape(), &[1, 3, 3, 3]);
///
/// // ReLU makes every output value non-negative.
/// for value in output.iter() {
///     assert!(*value >= 0.0);
/// }
/// ```
#[derive(Debug)]
pub struct DepthwiseConv2D {
    /// Number of input channels, which
    /// [`UnaryLayer::build`](crate::neural_network::traits::UnaryLayer::build) reads from the input
    /// shape
    channels: usize,
    /// Kernels per input channel. The output carries `channels * depth_multiplier` of them
    depth_multiplier: usize,
    /// Size of the convolution kernel as (height, width)
    kernel_size: (usize, usize),
    /// Stride of the convolution as (height_stride, width_stride)
    strides: (usize, usize),
    /// Tap spacing of the kernel as (height, width). 1 gives a solid kernel on that axis
    dilation_rate: (usize, usize),
    /// Padding strategy (Valid or Same)
    padding: PaddingType,
    /// 4D weight tensor with shape \[kernel_height, kernel_width, channels, depth_multiplier\]
    weights: Array4<f32>,
    /// 1D bias vector with shape \[channels * depth_multiplier\]
    ///
    /// The array stays allocated when `use_bias` is false. The forward pass and weight rosters
    /// exclude it.
    bias: Array1<f32>,
    /// Activation applied to the convolution output
    activation: Activation,
    /// Shape the layer was built for, batch axis first. `None` before the build
    built: Option<Shape>,
    /// Seed of the weight draw, or `None` to take the global seed or entropy
    random_state: Option<u64>,
    /// Whether the layer adds a bias to the convolution output
    use_bias: bool,
}

impl DepthwiseConv2D {
    /// Creates a new DepthwiseConv2D layer
    ///
    /// # Parameters
    ///
    /// - `kernel_size` - Size of the convolution kernel as (height, width)
    /// - `strides` - Stride of the convolution as (height_stride, width_stride)
    /// - `activation` - Activation function applied to the output
    ///
    /// # Returns
    ///
    /// - `Result<Self, Error>` - A new unbuilt `DepthwiseConv2D` layer, or an error
    ///
    /// # Notes
    ///
    /// There is no `filters` argument. A depthwise convolution derives its output channel count
    /// from the input, as `channels * depth_multiplier`. `depth_multiplier` defaults to 1. Set it
    /// with [`DepthwiseConv2D::with_depth_multiplier`]. Padding defaults to
    /// [`PaddingType::Valid`].
    ///
    /// Choose
    /// [`PaddingType::Same`] with
    /// [`DepthwiseConv2D::with_padding`]. The kernel is solid by default. Space its taps out with
    /// [`DepthwiseConv2D::with_dilation_rate`].
    ///
    /// The layer seeds weights from the global seed or
    /// entropy by default. For reproducible initialization, set a seed with
    /// [`DepthwiseConv2D::with_random_state`]
    ///
    /// # Errors
    ///
    /// - `Error::InvalidParameter` - If any kernel dimension or stride is 0
    /// - `Error::InvalidParameter` - If the activation carries an unusable parameter (see
    ///   [`Activation::validate`])
    pub fn new(
        kernel_size: (usize, usize),
        strides: (usize, usize),
        activation: impl Into<Activation>,
    ) -> Result<Self, Error> {
        validate_kernel_size_2d(kernel_size)?;
        validate_strides_2d(strides)?;
        let activation = activation.into();
        activation.validate()?;

        Ok(Self {
            channels: 0,
            depth_multiplier: 1,
            kernel_size,
            strides,
            dilation_rate: (1, 1),
            padding: PaddingType::Valid,
            weights: Array4::zeros((0, 0, 0, 0)),
            bias: Array1::zeros(0),
            activation,
            built: None,
            random_state: None,
            use_bias: true,
        })
    }

    /// Sets the padding mode (defaults to
    /// [`PaddingType::Valid`])
    ///
    /// # Parameters
    ///
    /// - `padding` - Padding strategy (Valid or Same)
    ///
    /// # Returns
    ///
    /// - `Self` - The updated layer
    pub fn with_padding(mut self, padding: PaddingType) -> Self {
        self.padding = padding;
        self
    }

    /// Sets the tap spacing of the kernel (defaults to `(1, 1)`)
    ///
    /// A dilation of `d` on an axis spaces the kernel taps `d` cells apart, so `k` taps span
    /// `(k - 1) * d + 1` input cells of that axis. The window still advances by the stride. A
    /// dilation of 1 on both axes gives a solid kernel
    ///
    /// # Parameters
    ///
    /// - `dilation_rate` - Tap spacing as (height, width)
    ///
    /// # Returns
    ///
    /// - `Result<Self, Error>` - The updated layer, or an error
    ///
    /// # Notes
    ///
    /// A depthwise convolution takes a stride above 1 and a dilation above 1 together. Only the
    /// plain and the transposed convolutions reject that pair
    ///
    /// The effective kernel is not bounded by the input axis here. Only
    /// [`PaddingType::Valid`] needs it to
    /// fit, and the build applies that rule
    ///
    /// # Errors
    ///
    /// - `Error::InvalidParameter` - If either dilation is 0
    pub fn with_dilation_rate(mut self, dilation_rate: (usize, usize)) -> Result<Self, Error> {
        let dilation = [dilation_rate.0, dilation_rate.1];
        validate_dilation(&dilation)?;
        self.dilation_rate = dilation_rate;
        Ok(self)
    }

    /// Sets how many kernels each input channel gets (defaults to 1)
    ///
    /// The layer then produces `channels * depth_multiplier` output channels. Input channel `c`'s
    /// multiplier `m` lands at output channel `c * depth_multiplier + m`. This re-runs weight
    /// initialization at the new shape and re-zeros the bias. Call it before you assign custom
    /// weights or start training
    ///
    /// # Parameters
    ///
    /// - `depth_multiplier` - Kernels per input channel
    ///
    /// # Returns
    ///
    /// - `Result<Self, Error>` - The updated layer, or an error
    ///
    /// # Errors
    ///
    /// - `Error::InvalidParameter` - If `depth_multiplier` is 0
    pub fn with_depth_multiplier(mut self, depth_multiplier: usize) -> Result<Self, Error> {
        validate_depth_multiplier(depth_multiplier)?;
        self.depth_multiplier = depth_multiplier;
        if self.built.is_some() {
            self.draw_parameters();
        }
        Ok(self)
    }

    /// Sets the depthwise weight seed and redraws weights if the layer has a build
    ///
    /// By default, the layer seeds weights from the global seed or entropy (see
    /// [`crate::random`]). For a built layer, this repeats Xavier/Glorot uniform initialization with `random_state`.
    /// Call it before you assign custom weights or start training. A built layer also resets
    /// its bias to 0
    ///
    /// # Parameters
    ///
    /// - `random_state` - Seed for weight initialization
    ///
    /// # Returns
    ///
    /// - `Self` - The updated layer
    pub fn with_random_state(mut self, random_state: u64) -> Self {
        self.random_state = Some(random_state);
        if self.built.is_some() {
            self.draw_parameters();
        }
        self
    }

    /// Xavier/Glorot uniform initialization of the \[kh, kw, channels, depth_multiplier\] tensor
    ///
    /// The depth multiplier replaces the filter count in the fan pair. Thus, `fan_in` counts
    /// every input channel although each depthwise unit reads only 1 channel. See
    /// [`Fans::conv`].
    fn draw_parameters(&mut self) {
        let (kernel_height, kernel_width) = self.kernel_size;
        let mut rng = crate::random::make_rng(self.random_state);
        // The depth multiplier takes the place of the filter count. See `Fans::conv`
        self.weights = Initializer::GlorotUniform.draw(
            (
                kernel_height,
                kernel_width,
                self.channels,
                self.depth_multiplier,
            ),
            Fans::conv(
                self.channels,
                self.depth_multiplier,
                kernel_height * kernel_width,
            ),
            &mut rng,
        );
        self.bias = Array1::zeros(self.channels * self.depth_multiplier);
    }

    /// Calculates padding dimensions for Same padding mode
    fn calculate_padding(
        &self,
        input_height: usize,
        input_width: usize,
        output_height: usize,
        output_width: usize,
    ) -> (usize, usize) {
        match self.padding {
            PaddingType::Valid => (0, 0),
            PaddingType::Same => {
                let (keff_h, keff_w) = self.effective_kernel_size();
                let pad_h =
                    ((output_height - 1) * self.strides.0 + keff_h).saturating_sub(input_height);
                let pad_w =
                    ((output_width - 1) * self.strides.1 + keff_w).saturating_sub(input_width);
                (pad_h, pad_w)
            }
        }
    }

    /// Sets whether the layer adds a bias to the convolution output (defaults to `true`)
    ///
    /// With `use_bias` set to false, the weight roster contains only the kernel. `param_count` counts it,
    /// and `parameters` yields it. A checkpoint holds 1 array under `<position>.kernel`. A
    /// checkpoint with bias cannot load into a layer without bias. The error names the path.
    ///
    /// # Parameters
    ///
    /// - `use_bias` - `true` to add a bias, `false` to leave it out
    ///
    /// # Returns
    ///
    /// - `Self` - The updated layer
    pub fn with_use_bias(mut self, use_bias: bool) -> Self {
        self.use_bias = use_bias;
        self
    }

    /// Sets the weights and bias for this layer
    ///
    /// # Parameters
    ///
    /// - `weights` - 4D weight tensor with shape
    ///   \[kernel_height, kernel_width, channels, depth_multiplier\]
    /// - `bias` - 1D bias vector with shape \[channels * depth_multiplier\], or `None`
    ///   for a layer built with [`with_use_bias(false)`](Self::with_use_bias)
    ///
    /// # Returns
    ///
    /// - `Result<(), Error>` - Success after the weights are set
    ///
    /// # Errors
    ///
    /// - `Error::NeuralNetwork(NnError::WeightShape)` - If `weights` or `bias` does not match
    ///   the existing shape
    /// - `Error::InvalidParameter` - If a bias is given to a layer that holds none, or none
    ///   is given to a layer that holds 1
    pub fn set_weights(
        &mut self,
        weights: Array4<f32>,
        bias: impl Into<Option<Array1<f32>>>,
    ) -> Result<(), Error> {
        if self.built.is_none() {
            return Err(Error::not_built("DepthwiseConv2D"));
        }
        validate_weight_shape("kernel", self.weights.shape(), weights.shape())?;
        let bias = validate_optional_weight("bias", "use_bias", self.use_bias, bias.into())?;
        if let Some(bias) = bias.as_ref() {
            validate_weight_shape("bias", self.bias.shape(), bias.shape())?;
        }
        self.weights = weights;
        if let Some(bias) = bias {
            self.bias = bias;
        }
        Ok(())
    }

    /// The layer's geometry for a given input, as the shared kernel wants it
    fn geometry(&self, input_shape: &[usize]) -> DepthwiseGeometry {
        let (height, width) = (input_shape[1], input_shape[2]);
        let output_shape = calculate_output_shape_2d(
            input_shape,
            self.effective_kernel_size(),
            self.strides,
            &self.padding,
        );
        let (out_height, out_width) = (output_shape[1], output_shape[2]);
        let (pad_h, pad_w) = self.calculate_padding(height, width, out_height, out_width);
        DepthwiseGeometry {
            input: (height, width),
            output: (out_height, out_width),
            channels: self.channels,
            depth_multiplier: self.depth_multiplier,
            kernel: self.kernel_size,
            strides: self.strides,
            dilation: self.dilation_rate,
            pad_before: (pad_h / 2, pad_w / 2),
        }
    }

    /// The extent the dilated taps span on each axis
    ///
    /// Every output-size and padding rule reads this rather than the kernel size. The 2 agree at
    /// a dilation of 1
    fn effective_kernel_size(&self) -> (usize, usize) {
        (
            effective_kernel(self.kernel_size.0, self.dilation_rate.0),
            effective_kernel(self.kernel_size.1, self.dilation_rate.1),
        )
    }

    /// Depthwise convolution over a channels-last tensor, with no activation
    ///
    /// The numeric body of [`UnaryLayer::forward`]. The result is the pre-activation `z`.
    /// [`UnaryLayer::forward`] applies the activation to it
    fn convolve(&self, input: &Tensor) -> Result<Tensor, Error> {
        validate_built_input(&self.built, "DepthwiseConv2D", input.shape())?;
        validate_valid_kernel_fits(
            self.padding.into(),
            &[self.kernel_size.0, self.kernel_size.1],
            &[self.dilation_rate.0, self.dilation_rate.1],
            &input.shape()[1..3],
        )?;

        let g = self.geometry(input.shape());
        let batch_size = input.shape()[0];
        let out_channels = g.out_channels();

        let input_std = input.as_standard_layout();
        let src = input_std
            .as_slice()
            .expect("standard-layout array is contiguous");
        let ker = self.weights.as_slice().expect("weights must be contiguous");
        let bias = self
            .use_bias
            .then(|| self.bias.as_slice().expect("bias must be contiguous"));

        let mut output = Array4::<f32>::zeros((batch_size, g.output.0, g.output.1, out_channels));
        depthwise_forward(
            &g,
            src,
            ker,
            bias,
            output.as_slice_mut().expect("output is contiguous"),
        );

        Ok(output.into_dyn())
    }
}

/// What the forward pass of [`DepthwiseConv2D`] parks for its backward pass
struct DepthwiseConv2DCache {
    /// The tensor the forward pass received
    input: Tensor,
    /// The tensor that the activation backward pass reads, as [`Activation::saves`] names
    activation: ActivationCache,
}

impl LayerBase for DepthwiseConv2D {
    fn layer_type(&self) -> &str {
        "DepthwiseConv2D"
    }

    fn param_count(&self) -> ParamCounts {
        // Read the arrays the layer holds rather than the configuration, so dropping the
        // bias corrects the count with no second formula to keep in step
        let bias = if self.use_bias { self.bias.len() } else { 0 };
        ParamCounts::trainable(self.weights.len() + bias)
    }

    fn parameters_mut(&mut self) -> Vec<ParamRef<'_>> {
        let Self {
            weights,
            bias,
            use_bias,
            ..
        } = self;
        let mut params = vec![ParamRef::weight(
            "kernel",
            weights.as_slice_mut().expect("weights must be contiguous"),
        )];
        if *use_bias {
            params.push(ParamRef::no_decay(
                "bias",
                bias.as_slice_mut().expect("bias must be contiguous"),
            ));
        }
        params
    }

    built_layer_shape_functions!();

    named_weight_layer_functions!(
        trainable "kernel" => weights,
        trainable "bias" => bias if use_bias,
    );
}

impl UnaryLayer for DepthwiseConv2D {
    /// Reads the channel count from the input shape, and draws the kernel and the bias
    fn build(&mut self, input: &Shape) -> Result<(), Error> {
        let Some(built) = start_build(&self.built, "DepthwiseConv2D", input)? else {
            return Ok(());
        };
        built.check_rank("DepthwiseConv2D", 4)?;
        let (batch, tail) = built.split_batch("DepthwiseConv2D")?;
        // The family validators check a full extent list. The batch extent is not part of it
        let mut dims = vec![batch.unwrap_or(1)];
        dims.extend(tail);
        validate_input_shape_2d(&dims)?;
        // The shape algebra holds every geometry rule, so a bad stack is refused before any
        // weight is drawn
        self.compute_output_shape(&built)?;
        self.channels = dims[3];
        self.built = Some(built);
        self.draw_parameters();
        Ok(())
    }

    fn forward(&self, input: &Tensor, ctx: &mut Ctx) -> Result<Tensor, Error> {
        if !self.is_built() {
            return Err(Error::not_built("DepthwiseConv2D"));
        }
        let z = self.convolve(input)?;
        if !ctx.is_training() {
            return self.activation.forward(&z);
        }
        let (activated, activation) = self.activation.forward_train(z)?;
        // Park only after a successful forward pass, so a rejected input leaves no partial state
        ctx.push_cache(
            "DepthwiseConv2D",
            DepthwiseConv2DCache {
                input: input.clone(),
                activation,
            },
        );
        Ok(activated)
    }

    fn backward(&self, grad_output: &Tensor, ctx: &mut Ctx) -> Result<Tensor, Error> {
        let cache: DepthwiseConv2DCache = ctx.pop_cache("DepthwiseConv2D")?;
        let grad_upstream = self.activation.backward(&cache.activation, grad_output)?;

        let input = &cache.input;

        let batch_size = input.shape()[0];
        let g = self.geometry(input.shape());

        let input_std = input.as_standard_layout();
        let src = input_std
            .as_slice()
            .expect("standard-layout array is contiguous");
        let grad_std = grad_upstream.as_standard_layout();
        let grad = grad_std
            .as_slice()
            .expect("standard-layout array is contiguous");
        let ker = self.weights.as_slice().expect("weights must be contiguous");

        let grads = depthwise_backward(&g, src, grad, ker, batch_size);

        ctx.add_grad(
            "kernel",
            Array4::from_shape_vec(self.weights.raw_dim(), grads.weight)
                .expect("weight gradient shape matches weights")
                .into_dyn(),
        )?;
        // A layer without bias has no stored bias gradient or bias optimizer state.
        if self.use_bias {
            ctx.add_grad("bias", Array1::from_vec(grads.bias).into_dyn())?;
        }

        Ok(
            Array4::from_shape_vec((batch_size, g.input.0, g.input.1, g.channels), grads.input)
                .expect("input gradient shape matches input")
                .into_dyn(),
        )
    }

    fn compute_output_shape(&self, input: &Shape) -> Result<Shape, Error> {
        input.check_rank("DepthwiseConv2D", 4)?;
        let (batch, tail) = input.split_batch("DepthwiseConv2D")?;
        let (height, width) = (tail[0], tail[1]);
        let channels = tail[2];
        // The layer keeps 1 kernel group per channel of the shape it was built for, so a built
        // layer refuses any other channel count. An unbuilt layer holds no group yet, and the
        // answer needs the channel count of the argument alone
        if self.built.is_some() && channels != self.channels {
            return Err(Error::dimension_mismatch(self.channels, channels));
        }
        let (keff_h, keff_w) = self.effective_kernel_size();
        // Reject `Valid` padding if the effective kernel exceeds an input axis. This prevents
        // an empty output axis.
        let (out_height, out_width) = match self.padding {
            PaddingType::Valid => (
                valid_output_size("DepthwiseConv2D", "height", height, keff_h, self.strides.0)?,
                valid_output_size("DepthwiseConv2D", "width", width, keff_w, self.strides.1)?,
            ),
            PaddingType::Same => (
                height.div_ceil(self.strides.0),
                width.div_ceil(self.strides.1),
            ),
        };
        Ok(Shape::from_batch(
            batch,
            &[out_height, out_width, channels * self.depth_multiplier],
        ))
    }
}

/// Unit tests for `DepthwiseConv2D`
#[cfg(test)]
mod tests {
    use super::*;
    use crate::neural_network::layers::activation::linear::Linear;
    use ndarray::ArrayD;

    /// Channels never mix. The values match a hand-worked cross-correlation
    ///
    /// Channel 0 carries 1..9 against an all-ones kernel. Channel 1 carries all 1 against an
    /// all-twos kernel. Channel 1's output is therefore the same constant at every position,
    /// while channel 0's output varies. This only holds if the 2 kernels stay on their own
    /// channels
    #[test]
    fn depthwise_forward_keeps_channels_separate_hand_derived() {
        let mut layer = DepthwiseConv2D::new((2, 2), (1, 1), Linear::new()).unwrap();
        layer.build(&Shape::known(&[1, 3, 3, 2])).unwrap();
        // [kh, kw, C, dm] with dm == 1: channel 0 all 1, channel 1 all 2
        let weights =
            Array4::from_shape_vec((2, 2, 2, 1), vec![1.0, 2.0, 1.0, 2.0, 1.0, 2.0, 1.0, 2.0])
                .unwrap();
        layer.set_weights(weights, Array1::zeros(2)).unwrap();

        // [1, 3, 3, 2]: channel 0 is 1..9, channel 1 is all 1
        let input = ArrayD::from_shape_vec(
            ndarray::IxDyn(&[1, 3, 3, 2]),
            vec![
                1.0, 1.0, 2.0, 1.0, 3.0, 1.0, // row 0
                4.0, 1.0, 5.0, 1.0, 6.0, 1.0, // row 1
                7.0, 1.0, 8.0, 1.0, 9.0, 1.0, // row 2
            ],
        )
        .unwrap();

        let out = layer.forward(&input, &mut Ctx::inference()).unwrap();
        assert_eq!(out.shape(), &[1, 2, 2, 2]);
        // Channel 0: the 4 window sums from the 2x2 kernel. Channel 1: 2 * 4 ones at every position
        assert_eq!(
            out.iter().copied().collect::<Vec<f32>>(),
            vec![12.0, 8.0, 16.0, 8.0, 24.0, 8.0, 28.0, 8.0]
        );
    }

    /// With `depth_multiplier`, input channel `c`'s multiplier `m` lands at output channel
    /// `c * depth_multiplier + m`
    #[test]
    fn depthwise_depth_multiplier_output_channel_order() {
        let mut layer = DepthwiseConv2D::new((1, 1), (1, 1), Linear::new())
            .unwrap()
            .with_depth_multiplier(2)
            .unwrap();
        layer.build(&Shape::known(&[1, 1, 1, 2])).unwrap();
        assert_eq!(layer.weights.shape(), &[1, 1, 2, 2]);

        // [kh=1, kw=1, C=2, dm=2] flat as (c, m): c0 -> [1, 10], c1 -> [100, 1000]
        let weights = Array4::from_shape_vec((1, 1, 2, 2), vec![1.0, 10.0, 100.0, 1000.0]).unwrap();
        layer.set_weights(weights, Array1::zeros(4)).unwrap();

        // 1 position holding [2, 3]
        let input = ArrayD::from_shape_vec(ndarray::IxDyn(&[1, 1, 1, 2]), vec![2.0, 3.0]).unwrap();

        let out = layer.forward(&input, &mut Ctx::inference()).unwrap();
        assert_eq!(out.shape(), &[1, 1, 1, 4]);
        // [c0m0, c0m1, c1m0, c1m1] = [2*1, 2*10, 3*100, 3*1000]
        assert_eq!(
            out.iter().copied().collect::<Vec<f32>>(),
            vec![2.0, 20.0, 300.0, 3000.0]
        );
    }
}