distdl.nn

Overview

These are the public interfaces.

Distributed Containers

Base Distributed Module

Base class of all distributed layers in DistDL.

Primitive Distributed Data Movement Layers

We implement a number of primitives.

Broadcast Layer	Performs a broadcast of a tensor from one partition to another.
Halo Exchange Layer	Performs the halo exchange operation of a tensor on a partition.
SumReduce Layer	Performs a sum-reduction of a tensor from one partition to another.
Transpose Layer	Performs a transpose or a tensor from one partition to another.

Distributed Comptutation Layers

We implement a number of distributed layers based on the actual layers and the primitives.

Convolution Layers	Distributed convolutional layers.
Pooling Layers	Distributed pooling layers.
Linear Layer	Distributed linear layer.
code_reference/nn/upsampling:Upsample Layer	Distributed upsampling layer.

Additional Sequential Layers

We implement some useful sequential modules.

code_reference/nn/interpolate:Interpolate Layer	N-dimensional, interpolation.
Padding Layer	N-dimensional, unbalanced padding.
Unpadding Layer	N-dimensional, unbalanced unpadding.