Add enum module support

[laggui]

Checklist

• Confirmed that run-checks all script has been executed.

Related Issues/PRs

Issue #726

#[derive(Module, Debug)]
enum ModuleEnum<B: Backend> {
    Basic(ModuleBasic<B>),
    Composed(ModuleComposed<B>),
}

Changes

Added support for #[derive(Module)] with enums.

• Added enum record derive support
• Added enum module derive support

In both cases, the generated methods defer the calls to the concrete types.

For example:

fn num_params(&self) -> usize {
    match self {
        Self::Basic(module) => burn::module::Module::<B>::num_params(module),
        Self::Composed(module) => burn::module::Module::<B>::num_params(module),
    }
}

fn load_record(self, record: Self::Record) -> Self {
    match self {
        Self::Basic(module) => {
            let Self::Record::Basic(r) = record else {panic!("Can't parse record from a different variant");};
            Self::Basic(burn::module::Module::<B>::load_record(module, r))
        }
        Self::Composed(module) => {
            let Self::Record::Composed(r) = record else {panic!("Can't parse record from a different variant");};
            Self::Composed(burn::module::Module::<B>::load_record(module, r))
        }
    }
}

Testing

Added enum module definitions in derive module tests and transposed some existing unit tests for enum modules.

[laggui]

Should we add something in the book to state that we support enum modules?

[codecov]

Codecov Report

Attention: 13 lines in your changes are missing coverage. Please review.

Comparison is base (4427768) 78.77% compared to head (a8ed24a) 78.87%.

Files	Patch %	Lines
crates/burn-core/tests/derive_module.rs	88.09%	5 Missing ⚠️
crates/burn-derive/src/record/item/codegen_enum.rs	93.82%	5 Missing ⚠️
crates/burn-derive/src/shared/enum_variant.rs	88.23%	2 Missing ⚠️
crates/burn-derive/src/record/base.rs	75.00%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1337      +/-   ##
==========================================
+ Coverage   78.77%   78.87%   +0.09%     
==========================================
  Files         551      555       +4     
  Lines       61836    62176     +340     
==========================================
+ Hits        48712    49040     +328     
- Misses      13124    13136      +12     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[laggui]

Now that I think about it, for simpler cases like num_params, visit and collect_devices which return a type that does not depend on the type we could generate a match with a single arm to capture all variants.

For example, num_params illustrated in the OP would look like this instead:

fn num_params(&self) -> usize {
    match self {
        Self::Basic(module) | Self::Composed(module) => burn::module::Module::<B>::num_params(module),
    }
}

Lmk if you have a preference.

[nathanielsimard]

LGTM, I think we can open issues to support all enum types, but for now it's not limiting, you can safely extract anything into modules and use a basic enum.

[nathanielsimard]

Is this supported?

#[derive(Module, Debug)]
enum ModuleEnumNamed<B: Backend> {
    Variant {
       fc1: nn::Linear<B>,
       fc2: nn::Linear<B>,
    },
}

If not maybe we can open an issue and support this in another PR.

[laggui]

Not yet! Your follow-up comment was correct. I only added support for unnamed fields/variants. If you try you should get a compile-time error (I added a panic to check for that specifically).

[nathanielsimard]

Looking at this, I think we don't yet support named enum, this can be added later.

[laggui]

Yeah only unnamed enum support in this PR. We could definitely add an issue and improve the support, I don't think it would require that much more work now that the table is set.