Module derive relies on B generics for Backend

[antimora]

Describe the bug

Getting this error:

error: the trait bound `Model<A>: Clone` is not satisfied
label: the trait `Clone` is not implemented for `Model<A>`

To Reproduce

Declare the following:

use burn::{
    module::Module,
    nn::{LayerNorm, Linear},
    tensor::backend::Backend,
};

// Define the model structure
#[derive(Module, Debug)]
pub struct Model<A: Backend> {
    layer_norm: LayerNorm<A>,
    output: Linear<A>,
}

Expected behavior

Should work just like for B: Backend

use burn::{
    module::Module,
    nn::{LayerNorm, Linear},
    tensor::backend::Backend,
};

// Define the model structure
#[derive(Module, Debug)]
pub struct Model<B: Backend> {
    layer_norm: LayerNorm<B>,
    output: Linear<B>,
}

[towerpark]

@nathanielsimard

Would it be all right if I look into this issue?
I'm not sure if it's fine to do so, since it has already been assigned.

[nathanielsimard]

You can work on this if you want!

[towerpark]

Thanks. I will try.

[towerpark]

I would like to hear your thoughts about my solution before starting to write code.

I'm thinking about adding a helper attribute for letting users provide backend information. An auto-detecting approach seems impossible because it would be too unreliable without the help of type information, e.g., checking if the name of a bound equals Backend won't work with user-defined backend traits or aliases such as use Backend as Be.

The syntax of the proposed approach is:

// Pattern 1:
//
// The target type has no generic type that implements a backend trait, and add one named
// `TYPE_NAME`.
//
// The `type` key is optional and defaults to `B`. Giving users a way to specify a name other
// than `B` can handle the case that the target type already has a generic type named `B`.
#[backend(add(type = "TYPE_NAME"))]

// Pattern 2:
//
// The target type already has a generic type named `TYPE_NAME` that implements a
// backend trait named `TRAIT_NAME`.
//
// The `type` key is optional and defaults to `B`.
//
// The `trait` key is optional and defaults to the first trait bound of `TYPE_NAME`. This key
// is added for handling the case that the backend trait is not the first one in the bounds.
#[backend(annotate(type = "TYPE_NAME", trait = "TRAIT_NAME"))]

// Pattern 3:
//
// The attribute is absent. Use the existing logic to check backend trait, i.e.:
//     if a generic type named `B` exists:
//         Go to Pattern 2 where `TYPE_NAME` is `B` and `TRAIT_NAME` is the name of the
//         first trait bound of `B`.
//     else:
//         Go to Pattern 1 where `TYPE_NAME` is `B`.
// This way we can ensure no existing code will be broken.

I'm also thinking about adding static assertions with no runtime overhead for verifying if the Backend trait is implemented to catch mis-annotations. The logic is as follows:

if it's Pattern 1:
    Verify that all generic types are not implementing the `Backend` trait.
else:
    Verify that `TRAIT_NAME` is the `Backend` trait itself or a subtrait of it.

The not-implementing assertion can be done by the approach suggested by this thread. I haven't dug into how to verify being some trait or being a subtrait, but I feel it's doable.

Plus, here are some note of my design choices:

• Why not merge the new attribute into the module helper attribute?
  Because the backend information is shared by Record and Module macros, and the module attribute is specific to Module.

• Why use the same syntax for both Module and Record macro while Record doesn't need the trait key?
  Reduces users' cognitive load by providing a unified syntax;
  Catches the case that a user thinks there's a generic type that implements the Backend trait but there actually isn't;
  The new attribute probably won't be used so often that the little verbosity matters.

• Why choose the path(key = value) form over path(value) for Pattern 1?
  Supports adding more keys in the future and makes its meaning clear.

• The default values of type and trait keys are chosen so as to be consistent with the existing behavior.

[nathanielsimard]

Hmm, not easy to answer, but I would say that the macro should be the same as the module. Maybe something like:

#[derive(Module)]
#[module(backend = "C")]
pub struct MyModule<C: Backend> {
   fc1: Linear<C>,
}

[towerpark]

I am sorry for the late reply.

I have several questions and would greatly appreciate it if you could clarify a little more.

• Do you mean the backend information should be placed solely in the module helper attribute and shared between the Module and Record macros? (because the Record macro also has the issue of hardcoded backend type)

• By backend = "C", do you mean using the existing logic of detecting and adding backend generic, but with a name specified by users instead of the hardcoded B?

• What about the hardcoded backend trait bound when a backend generic already exists?
  Adding another key-value pair to the module attribute is an option (e.g., #[module(backend = "C", backend_trait = "MyBackend")]), but I think grouping them together in some way would be easier to understand. I'm asking this because if we provide a way to customize the backend type, users may want a similar option for the backend trait too.