Repository: tracel-ai/burn
Pull request #2944: Check channels_in matches in convolution layers
Author: login: chlobes
Created: 2025-03-22T13:16:18Z | Merged: 2025-03-31T12:19:19Z

- [x] Confirmed that `run-checks all` script has been executed.
- [ ] Made sure the book is up to date with changes in this PR.

When channels_in to a convolution call mismatches the weights, backends break in various confusing ways.
This check avoids that with an early panic and informative error message.
I also added unit tests to ensure this check does panic.


Diff:
diff --git a/crates/burn-core/src/nn/conv/conv1d.rs b/crates/burn-core/src/nn/conv/conv1d.rs
index 9749e2873d7..4f4c9949edc 100644
--- a/crates/burn-core/src/nn/conv/conv1d.rs
+++ b/crates/burn-core/src/nn/conv/conv1d.rs
@@ -137,7 +137,7 @@ impl<B: Backend> Conv1d<B> {
     /// - input: `[batch_size, channels_in, length_in]`
     /// - output: `[batch_size, channels_out, length_out]`
     pub fn forward(&self, input: Tensor<B, 3>) -> Tensor<B, 3> {
-        let [_batch_size, _channels, length] = input.dims();
+        let length = input.dims()[2];
         let padding = self
             .padding
             .calculate_padding_1d(length, self.kernel_size, self.stride);
@@ -200,4 +200,14 @@ mod tests {
             "Conv1d {stride: 1, kernel_size: 5, dilation: 1, groups: 1, padding: Valid, params: 130}"
         );
     }
+
+    #[test]
+    #[should_panic = "Number of channels in input tensor and input channels of convolution must be equal. got: 4, expected: 5"]
+    fn input_channels_mismatch() {
+        let config = Conv1dConfig::new(5, 3, 3);
+        let conv = config.init::<TestBackend>(&Default::default());
+
+        let input = Tensor::<TestBackend, 3>::zeros([1, 4, 10], &Default::default());
+        let _ = conv.forward(input);
+    }
 }
diff --git a/crates/burn-core/src/nn/conv/conv2d.rs b/crates/burn-core/src/nn/conv/conv2d.rs
index 4ba929fbdd3..204568d0814 100644
--- a/crates/burn-core/src/nn/conv/conv2d.rs
+++ b/crates/burn-core/src/nn/conv/conv2d.rs
@@ -253,4 +253,14 @@ mod tests {
             "Conv2d {stride: [1, 1], kernel_size: [5, 5], dilation: [1, 1], groups: 1, padding: Valid, params: 126}"
         );
     }
+
+    #[test]
+    #[should_panic = "Number of channels in input tensor and input channels of convolution must be equal. got: 4, expected: 5"]
+    fn input_channels_mismatch() {
+        let config = Conv2dConfig::new([5, 3], [3, 3]);
+        let conv = config.init::<TestBackend>(&Default::default());
+
+        let input = Tensor::<TestBackend, 4>::zeros([1, 4, 10, 10], &Default::default());
+        let _ = conv.forward(input);
+    }
 }
diff --git a/crates/burn-core/src/nn/conv/conv3d.rs b/crates/burn-core/src/nn/conv/conv3d.rs
index 850fe97776e..57775835a39 100644
--- a/crates/burn-core/src/nn/conv/conv3d.rs
+++ b/crates/burn-core/src/nn/conv/conv3d.rs
@@ -249,4 +249,14 @@ mod tests {
             "Conv3d {stride: [1, 1, 1], kernel_size: [5, 5, 5], dilation: [1, 1, 1], groups: 1, padding: Valid, params: 626}"
         );
     }
+
+    #[test]
+    #[should_panic = "Number of channels in input tensor and input channels of convolution must be equal. got: 4, expected: 5"]
+    fn input_channels_mismatch() {
+        let config = Conv3dConfig::new([5, 3], [3, 3, 3]);
+        let conv = config.init::<TestBackend>(&Default::default());
+
+        let input = Tensor::<TestBackend, 5>::zeros([1, 4, 10, 10, 10], &Default::default());
+        let _ = conv.forward(input);
+    }
 }
diff --git a/crates/burn-core/src/nn/conv/conv_transpose1d.rs b/crates/burn-core/src/nn/conv/conv_transpose1d.rs
index e0b06bee047..bd90fc3c1a0 100644
--- a/crates/burn-core/src/nn/conv/conv_transpose1d.rs
+++ b/crates/burn-core/src/nn/conv/conv_transpose1d.rs
@@ -196,4 +196,14 @@ mod tests {
             "ConvTranspose1d {channels: [5, 2], stride: 1, kernel_size: 5, dilation: 1, groups: 1, padding: 0, padding_out: 0, params: 52}"
         );
     }
+
+    #[test]
+    #[should_panic = "Number of channels in input tensor and input channels of convolution must be equal. got: 4, expected: 5"]
+    fn input_channels_mismatch() {
+        let config = ConvTranspose1dConfig::new([5, 3], 3);
+        let conv = config.init::<TestBackend>(&Default::default());
+
+        let input = Tensor::<TestBackend, 3>::zeros([1, 4, 10], &Default::default());
+        let _ = conv.forward(input);
+    }
 }
diff --git a/crates/burn-core/src/nn/conv/conv_transpose2d.rs b/crates/burn-core/src/nn/conv/conv_transpose2d.rs
index 2d5695e973c..358dbcb3557 100644
--- a/crates/burn-core/src/nn/conv/conv_transpose2d.rs
+++ b/crates/burn-core/src/nn/conv/conv_transpose2d.rs
@@ -198,4 +198,14 @@ mod tests {
             "ConvTranspose2d {channels: [5, 2], stride: [1, 1], kernel_size: [5, 5], dilation: [1, 1], groups: 1, padding: [0, 0], padding_out: [0, 0], params: 252}"
         );
     }
+
+    #[test]
+    #[should_panic = "Number of channels in input tensor and input channels of convolution must be equal. got: 4, expected: 5"]
+    fn input_channels_mismatch() {
+        let config = ConvTranspose2dConfig::new([5, 3], [3, 3]);
+        let conv = config.init::<TestBackend>(&Default::default());
+
+        let input = Tensor::<TestBackend, 4>::zeros([1, 4, 10, 10], &Default::default());
+        let _ = conv.forward(input);
+    }
 }
diff --git a/crates/burn-core/src/nn/conv/conv_transpose3d.rs b/crates/burn-core/src/nn/conv/conv_transpose3d.rs
index d847a878476..a962894f356 100644
--- a/crates/burn-core/src/nn/conv/conv_transpose3d.rs
+++ b/crates/burn-core/src/nn/conv/conv_transpose3d.rs
@@ -202,4 +202,14 @@ mod tests {
             "ConvTranspose3d {channels: [5, 2], stride: [1, 1, 1], kernel_size: [5, 5, 5], dilation: [1, 1, 1], groups: 1, padding: [0, 0, 0], padding_out: [0, 0, 0], params: 1252}"
         );
     }
+
+    #[test]
+    #[should_panic = "Number of channels in input tensor and input channels of convolution must be equal. got: 4, expected: 5"]
+    fn input_channels_mismatch() {
+        let config = ConvTranspose3dConfig::new([5, 3], [3, 3, 3]);
+        let conv = config.init::<TestBackend>(&Default::default());
+
+        let input = Tensor::<TestBackend, 5>::zeros([1, 4, 10, 10, 10], &Default::default());
+        let _ = conv.forward(input);
+    }
 }
diff --git a/crates/burn-core/src/nn/conv/deform_conv2d.rs b/crates/burn-core/src/nn/conv/deform_conv2d.rs
index 26d8f553c26..476d5db6312 100644
--- a/crates/burn-core/src/nn/conv/deform_conv2d.rs
+++ b/crates/burn-core/src/nn/conv/deform_conv2d.rs
@@ -275,4 +275,15 @@ mod tests {
             "DeformConv2d {stride: [1, 1], kernel_size: [5, 5], dilation: [1, 1], weight_groups: 1, offset_groups: 1, padding: Valid, params: 126}"
         );
     }
+
+    #[test]
+    #[should_panic = "Number of channels in input tensor and input channels of convolution must be equal. got: 4, expected: 5"]
+    fn input_channels_mismatch() {
+        let config = DeformConv2dConfig::new([5, 3], [3, 3]);
+        let conv = config.init::<TestBackend>(&Default::default());
+
+        let input = Tensor::<TestBackend, 4>::zeros([1, 4, 10, 10], &Default::default());
+        let offset = Tensor::<TestBackend, 4>::zeros([1, 2 * 3 * 3, 10, 10], &Default::default());
+        let _ = conv.forward(input, offset, None);
+    }
 }
diff --git a/crates/burn-tensor/src/tensor/api/check.rs b/crates/burn-tensor/src/tensor/api/check.rs
index 16f91f57d0d..e4fbd516ffc 100644
--- a/crates/burn-tensor/src/tensor/api/check.rs
+++ b/crates/burn-tensor/src/tensor/api/check.rs
@@ -1194,6 +1194,45 @@ impl TensorCheck {
 
         check
     }
+
+    /// Checks if input is compatible with convolution weights.
+    pub fn conv<const D1: usize, const D2: usize>(
+        ops: &str,
+        x: [usize; D1],
+        weight: [usize; D2],
+        groups: usize,
+    ) -> Self {
+        let mut check = TensorCheck::Ok;
+        let channels = x[1];
+        let expected = weight[1] * groups;
+        if channels != expected {
+            check = check.register(
+                ops,
+                TensorError::new("Number of channels in input tensor and input channels of convolution must be equal.")
+                .details(format!("got: {channels}, expected: {expected}")),
+            );
+        }
+        check
+    }
+
+    /// Checks if input is compatible with transposed convolution weights.
+    pub fn conv_transpose<const D1: usize, const D2: usize>(
+        ops: &str,
+        x: [usize; D1],
+        weight: [usize; D2],
+    ) -> Self {
+        let mut check = TensorCheck::Ok;
+        let channels = x[1];
+        let expected = weight[0];
+        if channels != expected {
+            check = check.register(
+                ops,
+                TensorError::new("Number of channels in input tensor and input channels of convolution must be equal.")
+                .details(format!("got: {channels}, expected: {expected}")),
+            );
+        }
+        check
+    }
 }
 
 pub(crate) struct FailedTensorCheck {
diff --git a/crates/burn-tensor/src/tensor/module.rs b/crates/burn-tensor/src/tensor/module.rs
index 1c9c8b70455..1d7afaaa2e6 100644
--- a/crates/burn-tensor/src/tensor/module.rs
+++ b/crates/burn-tensor/src/tensor/module.rs
@@ -1,6 +1,8 @@
 use crate::{
     Int, Tensor, TensorPrimitive,
     backend::Backend,
+    check,
+    check::TensorCheck,
     ops::{ConvOptions, ConvTransposeOptions, InterpolateOptions, UnfoldOptions},
 };
 
@@ -27,6 +29,12 @@ pub fn conv1d<B>(
 where
     B: Backend,
 {
+    check!(TensorCheck::conv(
+        "conv1d",
+        x.dims(),
+        weight.dims(),
+        options.groups,
+    ));
     Tensor::new(TensorPrimitive::Float(B::conv1d(
         x.primitive.tensor(),
         weight.primitive.tensor(),
@@ -45,6 +53,12 @@ pub fn conv2d<B>(
 where
     B: Backend,
 {
+    check!(TensorCheck::conv(
+        "conv2d",
+        x.dims(),
+        weight.dims(),
+        options.groups,
+    ));
     Tensor::new(TensorPrimitive::Float(B::conv2d(
         x.primitive.tensor(),
         weight.primitive.tensor(),
@@ -63,6 +77,12 @@ pub fn conv3d<B>(
 where
     B: Backend,
 {
+    check!(TensorCheck::conv(
+        "conv3d",
+        x.dims(),
+        weight.dims(),
+        options.groups,
+    ));
     Tensor::new(TensorPrimitive::Float(B::conv3d(
         x.primitive.tensor(),
         weight.primitive.tensor(),
@@ -83,6 +103,12 @@ pub fn deform_conv2d<B>(
 where
     B: Backend,
 {
+    check!(TensorCheck::conv(
+        "deform_conv2d",
+        x.dims(),
+        weight.dims(),
+        options.weight_groups,
+    ));
     Tensor::new(TensorPrimitive::Float(B::deform_conv2d(
         x.primitive.tensor(),
         offset.primitive.tensor(),
@@ -103,6 +129,11 @@ pub fn conv_transpose1d<B>(
 where
     B: Backend,
 {
+    check!(TensorCheck::conv_transpose(
+        "conv_transpose1d",
+        x.dims(),
+        weight.dims(),
+    ));
     Tensor::new(TensorPrimitive::Float(B::conv_transpose1d(
         x.primitive.tensor(),
         weight.primitive.tensor(),
@@ -121,6 +152,11 @@ pub fn conv_transpose2d<B>(
 where
     B: Backend,
 {
+    check!(TensorCheck::conv_transpose(
+        "conv_transpose2d",
+        x.dims(),
+        weight.dims(),
+    ));
     Tensor::new(TensorPrimitive::Float(B::conv_transpose2d(
         x.primitive.tensor(),
         weight.primitive.tensor(),
@@ -139,6 +175,11 @@ pub fn conv_transpose3d<B>(
 where
     B: Backend,
 {
+    check!(TensorCheck::conv_transpose(
+        "conv_transpose3d",
+        x.dims(),
+        weight.dims(),
+    ));
     Tensor::new(TensorPrimitive::Float(B::conv_transpose3d(
         x.primitive.tensor(),
         weight.primitive.tensor(),


Reviews:
user: login: laggui
body: Thanks for the added validation!

If we are going to perform a check, I think it would be at the function level in `burn-tensor/src/tensor/modules.rs` 🤔 so it is always checked even when using the functional form.

Also, what do you think about a debug assertion for such runtime checks? During development these can be enabled, and then disabled in release mode.
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-25T13:32:37Z
user: login: laggui
body: Sounds good!

We can simplify this a bit, then should be good to go 🙂 
state: CHANGES_REQUESTED
author_association: MEMBER
submitted_at: 2025-03-27T12:54:54Z
user: login: chlobes
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-28T00:00:11Z
user: login: chlobes
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-28T00:06:10Z
user: login: laggui
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-28T12:44:42Z
user: login: laggui
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-28T12:56:13Z
user: login: laggui
body: > Is this good?

Based on our discussion it made more sense to split the conv vs transpose conv checks made sense. Made the changes myself since you had already split it initially 😅 Also removed an accidentally committed file.

LGTM now! But lmk if you have any other comments, will wait before merging.
state: APPROVED
author_association: MEMBER
submitted_at: 2025-03-28T13:17:25Z

Review comments:
diff_hunk: @@ -137,7 +137,13 @@ impl<B: Backend> Conv1d<B> {
     /// - input: `[batch_size, channels_in, length_in]`
     /// - output: `[batch_size, channels_out, length_out]`
     pub fn forward(&self, input: Tensor<B, 3>) -> Tensor<B, 3> {
-        let [_batch_size, _channels, length] = input.dims();
+        let [_batch_size, channels_in, length] = input.dims();
+        let expected = self.weight.dims()[1] * self.groups;
+        assert_eq!(
+            channels_in, expected,
+            "This conv layer requies a channels_in dimension of {expected}, but got {channels_in}"
path: crates/burn-core/src/nn/conv/conv1d.rs
user: login: laggui
body: Typo: `requies` -> `requires`.

I would also simplify the message to: `"Expected input to have {expected} channels, but got {channels_in} instead"`
created_at: 2025-03-25T13:20:14Z
author_association: MEMBER
diff_hunk: @@ -1194,6 +1195,62 @@ impl TensorCheck {
 
         check
     }
+
+    pub fn conv<const D1: usize, const D2: usize, const N: usize>(
+        ops: &str,
+        x: [usize; D1],
+        weight: [usize; D2],
+        options: &ConvOptions<N>,
+    ) -> Self {
+        let mut check = TensorCheck::Ok;
+        let channels = x[1];
+        let expected = weight[1] * options.groups;
+        if channels != expected {
+            check = check.register(
+                ops,
+                TensorError::new("Number of channels in input tensor and input channels of convolution must be equal.")
+                .details(format!("got: {channels}, expected: {expected}")),
+            );
+        }
+        check
+    }
path: crates/burn-tensor/src/tensor/api/check.rs
user: login: laggui
body: The checks are pretty much the same for all convolution types. I think we can keep a single check then, and use `TensorCheck::conv(...)` with different `ops` value.

```rust
pub fn conv(
    ops: &str,
    input_shape: Shape,
    weight_shape: Shape,
    groups: usize,
) -> Self {
    let mut check = TensorCheck::Ok;
    let input_channels = input_shape.dims[1];
    let expected = weight_shape.dims[1] * groups;
    if channels != expected {
        check = check.register(
            ops,
            TensorError::new("Number of channels in input tensor must match the convolution input channels.")
            .details(format!("Got {channels} channels, expected {expected}")),
        );
    }
    check
}
```

Btw transpose conv options also have `groups`.
created_at: 2025-03-27T12:54:25Z
author_association: MEMBER
diff_hunk: @@ -1194,6 +1195,62 @@ impl TensorCheck {
 
         check
     }
+
+    pub fn conv<const D1: usize, const D2: usize, const N: usize>(
+        ops: &str,
+        x: [usize; D1],
+        weight: [usize; D2],
+        options: &ConvOptions<N>,
+    ) -> Self {
+        let mut check = TensorCheck::Ok;
+        let channels = x[1];
+        let expected = weight[1] * options.groups;
+        if channels != expected {
+            check = check.register(
+                ops,
+                TensorError::new("Number of channels in input tensor and input channels of convolution must be equal.")
+                .details(format!("got: {channels}, expected: {expected}")),
+            );
+        }
+        check
+    }
path: crates/burn-tensor/src/tensor/api/check.rs
user: login: chlobes
body: They do have groups, but (I think?) those groups are not relevant to channels_in size
created_at: 2025-03-28T00:00:11Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1194,6 +1195,62 @@ impl TensorCheck {
 
         check
     }
+
+    pub fn conv<const D1: usize, const D2: usize, const N: usize>(
+        ops: &str,
+        x: [usize; D1],
+        weight: [usize; D2],
+        options: &ConvOptions<N>,
+    ) -> Self {
+        let mut check = TensorCheck::Ok;
+        let channels = x[1];
+        let expected = weight[1] * options.groups;
+        if channels != expected {
+            check = check.register(
+                ops,
+                TensorError::new("Number of channels in input tensor and input channels of convolution must be equal.")
+                .details(format!("got: {channels}, expected: {expected}")),
+            );
+        }
+        check
+    }
path: crates/burn-tensor/src/tensor/api/check.rs
user: login: chlobes
body: Reasoning why I used `[usize; D]` over `Shape` is to make index out of bounds errors statically impossible.
In fact now I think about it, to really ensure that we should have a comptime assertion.

Reasoning for having seperate methods for each conv type was so that the conv function itself doesn't need to "know", but that's admittedly flismy reasoning and I'm happy to reduce codesize 🙂
created_at: 2025-03-28T00:06:09Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1194,6 +1195,62 @@ impl TensorCheck {
 
         check
     }
+
+    pub fn conv<const D1: usize, const D2: usize, const N: usize>(
+        ops: &str,
+        x: [usize; D1],
+        weight: [usize; D2],
+        options: &ConvOptions<N>,
+    ) -> Self {
+        let mut check = TensorCheck::Ok;
+        let channels = x[1];
+        let expected = weight[1] * options.groups;
+        if channels != expected {
+            check = check.register(
+                ops,
+                TensorError::new("Number of channels in input tensor and input channels of convolution must be equal.")
+                .details(format!("got: {channels}, expected: {expected}")),
+            );
+        }
+        check
+    }
path: crates/burn-tensor/src/tensor/api/check.rs
user: login: laggui
body: > They do have groups, but (I think?) those groups are not relevant to channels_in size

Ahh you're right, the groups affect the channels_out in conv transpose. Forgot about that for a moment
created_at: 2025-03-28T12:44:42Z
author_association: MEMBER
diff_hunk: @@ -1194,6 +1195,62 @@ impl TensorCheck {
 
         check
     }
+
+    pub fn conv<const D1: usize, const D2: usize, const N: usize>(
+        ops: &str,
+        x: [usize; D1],
+        weight: [usize; D2],
+        options: &ConvOptions<N>,
+    ) -> Self {
+        let mut check = TensorCheck::Ok;
+        let channels = x[1];
+        let expected = weight[1] * options.groups;
+        if channels != expected {
+            check = check.register(
+                ops,
+                TensorError::new("Number of channels in input tensor and input channels of convolution must be equal.")
+                .details(format!("got: {channels}, expected: {expected}")),
+            );
+        }
+        check
+    }
path: crates/burn-tensor/src/tensor/api/check.rs
user: login: laggui
body: > Reasoning why I used [usize; D] over Shape is to make index out of bounds errors statically impossible.

The checks are all used within conv functions that have fixed const rank `D`. So if the tensors are invalid (weight or input) it would already fail to compile at that level. Also, other such checks are currently at runtime (due to the const generics issue you highlight), so I don't think it's required.

I understand the motivation though!
created_at: 2025-03-28T12:56:13Z
author_association: MEMBER

PR comments:
user: login: codecov[bot]
created_at: 2025-03-23T14:27:45Z
body: ## [Codecov](https://app.codecov.io/gh/tracel-ai/burn/pull/2944?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=tracel-ai) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 82.17%. Comparing base [(`bf4b7e0`)](https://app.codecov.io/gh/tracel-ai/burn/commit/bf4b7e033e329a0bc5d7e016a56b21bb8f9094a3?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=tracel-ai) to head [(`e0af003`)](https://app.codecov.io/gh/tracel-ai/burn/commit/e0af0032dfd1ac1c01d9587a4fd55e549c3b2194?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=tracel-ai).
> Report is 17 commits behind head on main.

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##             main    #2944      +/-   ##
==========================================
+ Coverage   82.11%   82.17%   +0.05%     
==========================================
  Files         866      871       +5     
  Lines      120137   120511     +374     
==========================================
+ Hits        98656    99035     +379     
+ Misses      21481    21476       -5     
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/tracel-ai/burn/pull/2944?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=tracel-ai).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=tracel-ai).

<details><summary> :rocket: New features to boost your workflow: </summary>

- :snowflake: [Test Analytics](https://docs.codecov.com/docs/test-analytics): Detect flaky tests, report on failures, and find test suite problems.
</details>
author_association: NONE
user: login: chlobes
created_at: 2025-03-27T07:40:10Z
body: If we're doing the checks at the function level, then it seems sensible to use `check!`, I'll take a look at writing TensorCheck method(s) for that.

I'm not sure about runtime vs debug checks, I often build in release while debugging because debug gives me eg OOM or stack overflow errors.
 I think whatever solution is used could be put inside the `check!` macro so it gets applied uniformly.
 Could just feature flag the whole thing, `[cfg(not(feature = "elide_checks"))]` so that turning on `elide-checks`, would compile out all the checks.
 All this is obviously out of scope of course.
author_association: CONTRIBUTOR
user: login: chlobes
created_at: 2025-03-28T04:11:29Z
body: Is this good?
author_association: CONTRIBUTOR
user: login: chlobes
created_at: 2025-03-29T02:24:29Z
body: Thanks for your help, LGTM!
author_association: CONTRIBUTOR