Skip to content

sciml.layers

The layers and functions of the forward pass of a network.

Every layer and every function of the NN YAML is implemented once, against a backend, and registered under its PyTorch name in LAYERS or FUNCTIONS with the backends it supports. Importing this package registers all of them.

module content backends
core Linear, Bilinear, Flatten, the dropout layers numpy, sympy
functions the activation functions, flatten, cat numpy, sympy
convolution Conv1-3d, ConvTranspose1-3d numpy
pooling MaxPool, AvgPool, LPPool and the adaptive pools, 1-3d numpy
normalization BatchNorm1-3d, InstanceNorm1-3d, LayerNorm numpy

ArraySpec dataclass

ArraySpec(shape, required=True, trainable=True)

An array of a layer.

Attributes:

Name Type Description
shape tuple[int, ...]

the shape of the array in the PyTorch layout.

required bool

whether the layer cannot be evaluated without the array.

trainable bool

whether the elements are parameters of a fit. The running statistics of a normalization layer are arrays but not parameters.

FunctionType dataclass

FunctionType(name, function, backends)

A function or method of the forward pass, e.g. relu.

Attributes:

Name Type Description
name str

the name of the PyTorch function.

function Callable[..., ndarray]

the implementation (backend, *inputs, **kwargs).

backends frozenset[BackendKind]

the backends the implementation supports.

LayerType dataclass

LayerType(name, forward, arrays, backends)

A type of layer, e.g. Linear.

Attributes:

Name Type Description
name str

the name of the PyTorch class.

forward Callable[..., ndarray]

the implementation (backend, args, arrays, *inputs).

arrays ArraysFunction

the arrays of a layer of this type from its arguments.

backends frozenset[BackendKind]

the backends the implementation supports.