Struct DataParallel 

pub struct DataParallel<M> {
    pub module: M,
    pub num_workers: usize,
}

Wraps a Module and splits each input batch across num_workers threads.

The forward pass:

1. Split input along dimension 0 into num_workers chunks
2. Run each chunk through the module in parallel (rayon)
3. Concatenate the outputs

Because all workers share the same parameters (Tensor uses Arc), the autograd graph correctly tracks all operations. After computing loss on the concatenated output and calling .backward(), the gradients are automatically accumulated across all chunks.

Example

let model = Sequential::new(vec![...]);
let dp = DataParallel::new(model, 4);  // 4 workers
let output = dp.forward(&big_batch)?;  // splits into 4 chunks

Fields

module: M

The underlying module (shared across workers).

num_workers: usize

Number of parallel workers.

Implementations

impl<M> DataParallel<M>

pub fn new(module: M, num_workers: usize) -> Self

Create a DataParallel wrapper with the given number of workers.

num_workers controls how many chunks the batch is split into. For CPU, this maps to rayon thread-pool parallelism.

pub fn inner(&self) -> &M

Get a reference to the underlying module.

pub fn inner_mut(&mut self) -> &mut M

Get a mutable reference to the underlying module.

pub fn into_inner(self) -> M

Unwrap the DataParallel, returning the inner module.

Trait Implementations

impl<M: Clone> Clone for DataParallel<M>

fn clone(&self) -> Self

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl<M: Debug> Debug for DataParallel<M>

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.

impl<M, B> Module<B> for DataParallel<M>

where M: Module<B> + Send + Sync, B: Backend,

fn forward(&self, x: &Tensor<B>) -> Result<Tensor<B>>

Compute the output tensor from the input tensor. This defines the layer’s computation (forward pass).

fn parameters(&self) -> Vec<Tensor<B>>

Return all trainable parameters of this module. The optimizer uses these to update weights during training.

fn named_parameters(&self) -> Vec<(String, Tensor<B>)>

Return all trainable parameters with human-readable names.

fn set_training(&self, training: bool)

Set training or evaluation mode.

fn is_training(&self) -> bool

Whether the module is in training mode (default: true).

fn train(&self)

Convenience: set training mode.

fn eval(&self)

Convenience: set evaluation mode.

fn num_parameters(&self) -> usize

Total number of scalar parameters in this module.

fn trainable_params_count(&self) -> usize

Number of trainable (variable) parameters.

fn frozen_parameters(&self) -> Vec<Tensor<B>>

Freeze all parameters: returns new parameter tensors with is_variable = false, preventing gradient accumulation.

fn state_dict(&self) -> Vec<(String, Tensor<B>)>

Returns a state_dict-style map of parameter name → tensor.