huggingface / huggingface/candle

[Proposal] [candle_nn] Redesign Module/ModuleT traits

Open
#1,506 3 comments 0 reactions 0 assignees View on GitHub
Dominant language
Rust
Stars
21k
Forks
1.8k
Avg merge
16h 42m
Merged PRs (30d)
25

Description

While implementing #1504, I noticed that the the blanket implementation of ModuleT for Modules makes it awkward to create different behavior between training and learning. For a module to use the train parameter of the ModuleT trait, it cannot implement Module, because doing so will cause a "conflicting implementations" error at compile time. I believe that removing the ModuleT trait and adding the `forward_t` method to the Module trait can make this api easier to use.
```rust
trait Module {
fn forward(&self, xs: &Tensor) -> Result;
fn forward_t(&self, xs: &Tensor, _train: bool) -> Result {
self.forward(xs)
}
}
```
Note that the above implementation is nearly the same to the end user but with more flexibility. Implementations of `forward` only would look exactly the same, while implementations of `forward_t` only would look like the following:
```rust
impl Module for Mlp {
fn forward(&self, xs: &Tensor) -> Result {
self.forward_t(xs, false)
}

fn forward_t(&self, xs: &Tensor, train: bool) -> Result {
// ... long, complicated implementation passing 'train' to submodules
}
}
```
We could even make the implementation of `forward` above the default implementation, though this could lead to users doing `impl Module for Mlp {}` and creating infinite recursion.

These changes would not be difficult to implement, but I understand if they need to wait until a major/minor release.

Contributor guide

No contributing guide indexed for this repository

Research direction

Start by locating the candle_nn Module and ModuleT trait definitions and reviewing the behavior described alongside issue #1504. Trace their blanket implementations and existing forward calls to assess the API impact. Done means training-specific behavior can be implemented without conflicting implementations while existing forward-only implementations retain their behavior.

Written by the indexing model from the issue text.

Assessment

Tech stack
rust
Domain
machine-learning
Issue type
Refactor
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.