How to create modules that contain other generic modules?

[mentics]

I would like to create a module that can contain other modules that implement a trait. Is there a way to do that?

Something like the following, but it gives errors. I tried to add some associated type bounds but then the proc-macro panicked (unexpected token).

use burn::prelude::*;

trait HasForward<B: Backend> {
    fn forward<const D: usize>(&self, input: Tensor<B, D>) -> Tensor<B, D>;
}

#[derive(Module, Debug)]
pub struct AcrossDim<B: Backend, L: Module<B> + HasForward<B>> {
    modules: Vec<L>
}

Error:

error[E0277]: the trait bound `<L as burn::module::AutodiffModule<B>>::InnerModule: test::HasForward<<B as burn::tensor::backend::AutodiffBackend>::InnerBackend>` is not satisfied
 --> src/test.rs:7:10
  |
7 | #[derive(Module, Debug)]
  |          ^^^^^^ the trait `test::HasForward<<B as burn::tensor::backend::AutodiffBackend>::InnerBackend>` is not implemented for `<L as burn::module::AutodiffModule<B>>::InnerModule`
  |
note: required by a bound in `test::AcrossDim`
 --> src/test.rs:8:49
  |
8 | pub struct AcrossDim<B: Backend, L: Module<B> + HasForward<B>> {
  |                                                 ^^^^^^^^^^^^^ required by this bound in `AcrossDim`
  = note: this error originates in the derive macro `Module` (in Nightly builds, run with -Z macro-backtrace for more info)

[antimora]

@nathanielsimard would know of this limitation. Is this expected? Should we file an issue ticket for this?

[nathanielsimard]

This is supported, you just have to avoid specifying the trait bounds in the module definition, but rather in the implementation block:

use burn::prelude::*;
use core::marker::PhantomData;

trait HasForward<B: Backend> {
    fn forward<const D: usize>(&self, input: Tensor<B, D>) -> Tensor<B, D>;
}

#[derive(Module, Debug)]
pub struct AcrossDim<B: Backend, L> {
    modules: Vec<L>
    _backend: PhantomData<B>,
}

impl<B: Backend, L: HasForward<B>> AcrossDim<B, L> {
    // implementation goes here
}