kd.nn.NormModule

class kauldron.modules.NormModule(*args, **kwargs)

Bases: Protocol

Interface specification for norm modules (to be used as type annotation).

Module should be callable with a single tensor (x) of any shape and return a tensor of the same shape.

Example modules include: nn.LayerNorm, nn.RMSNorm, nn.GroupNorm, nn.BatchNorm.