ESCN Layers

class dipm.layers.escn.layernorm.LayerNormType(*values)

Options for the LayerNorm of the EquiformerV2 model.

dipm.layers.escn.layernorm.get_layernorm_layer(norm_type: LayerNormType | str, lmax: int, num_channels: int, eps: float = 1e-05, affine: bool = True, normalization: str = 'component', *, dtype: str | type[Any] | dtype | SupportsDType | Any | None = None, param_dtype: str | type[Any] | dtype | SupportsDType | Any = <class 'jax.numpy.float32'>, rngs: Rngs)