Skip to main content

burn_core/module/parallel/
layer.rs

1use crate::module::Module;
2
3/// One layer of a model split by [`LayerParallelism`](super::LayerParallelism). It runs on the
4/// device its placement gives it, and what it returns moves to the next layer's device.
5pub trait DistributedLayer: Module {
6    /// What the layer takes, already on its device.
7    type Input;
8    /// What the layer returns, left on its device.
9    type Output;
10
11    /// Run the layer.
12    fn forward(&self, input: Self::Input) -> Self::Output;
13}