Python module

layer_norm

Layer Normalization layer.

`ConstantLayerNorm`

class max.nn.norm.layer_norm.ConstantLayerNorm(dims, device, dtype, eps=1e-05)

Layer normalization block with constant gamma and beta values.

Parameters:

beta: npt.NDArray[np.floating[Any]]

device: DeviceRef

dtype: DType

eps: float = 1e-05

gamma: npt.NDArray[np.floating[Any]]

class max.nn.norm.layer_norm.LayerNorm(dims, devices, dtype, eps=1e-05, use_bias=True)

Layer normalization block.

Parameters:

shard(devices)

Creates sharded views of this LayerNorm across multiple devices.

Parameters:: devices (Iterable[DeviceRef]) – Iterable of devices to place the shards on.
Returns:: List of sharded LayerNorm instances, one for each device.
Return type:: Sequence[LayerNorm]

property sharding_strategy: ShardingStrategy | None

Get the LayerNorm sharding strategy.