mojolearn.LayerNorm

class mojolearn.LayerNorm(normalized_shape, eps=1e-05, elementwise_affine=True, bias=True, numeric_mode=None)

Bases: object

Layer normalization layer, as torch.nn.LayerNorm.

Holds a weight (initialized to ones) and a bias (zeros) over normalized_shape. forward(x) (also layer(x)) normalizes the trailing dimensions of x; backward(dy), called after a forward, returns dx and stores the parameter gradients in weight_grad and bias_grad. Inputs must be float32.

Parameters:

• normalized_shape (int or tuple of int) – Trailing shape to normalize over. Every dimension must be positive.

• eps (float, default 1e-5) – Added to the variance inside the square root.

• elementwise_affine (bool, default True) – Learn a weight (and a bias when bias is True).

• bias (bool, default True) – Learn a bias. Ignored when elementwise_affine is False.

• numeric_mode ({‘identical’, ‘deterministic’, ‘fast’} or None, default None) – Numeric mode. None uses the library default.

Variables:

• weight (ndarray or None) – Elementwise scale, None without elementwise_affine.

• bias (ndarray or None) – Elementwise shift, None without a bias.

• weight_grad (ndarray or None) – Gradient of weight from the last backward.

• bias_grad (ndarray or None) – Gradient of bias from the last backward.

Examples

>>> import numpy as np, mojolearn
>>> x = np.random.default_rng(0).random((32, 64), dtype=np.float32)
>>> ln = mojolearn.LayerNorm(64)
>>> y = ln(x)
>>> dx = ln.backward(np.ones_like(y))

__init__(normalized_shape, eps=1e-05, elementwise_affine=True, bias=True, numeric_mode=None)

Methods

backward(dy)	Backpropagate through the last forward.
forward(x)	Normalize x over the trailing normalized_shape dimensions.
to_device(a)	Upload a float32 array once so forward and backward keep it on the device.

to_device(a)

Upload a float32 array once so forward and backward keep it on the device.

Passing the result to forward keeps the output and gradients on the device and avoids a copy in each direction per call. When this install has no device path, the array itself is returned.

Parameters:

a (ndarray of float32) – Array to upload.

Returns:

device tensor or ndarray

forward(x)

Normalize x over the trailing normalized_shape dimensions.

Parameters:

x (ndarray of float32) – Input whose trailing shape is normalized_shape; otherwise ValueError.

Returns:

y (ndarray of float32) – Normalized output, the shape of x.

backward(dy)

Backpropagate through the last forward.

Sets weight_grad and bias_grad. Raises RuntimeError when no forward has run.

Parameters:

dy (ndarray of float32) – Gradient with respect to the output, the shape of the forward input.

Returns:

dx (ndarray of float32) – Gradient with respect to the forward input.