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: ~dipm.layers.escn.layernorm.LayerNormType | str, lmax: int, num_channels: int, eps: float = 1e-05, affine: bool = True, normalization: str = 'component', *, dtype: str | type[~typing.Any] | ~numpy.dtype | ~jax._src.typing.SupportsDType | ~typing.Any | None = None, param_dtype: str | type[~typing.Any] | ~numpy.dtype | ~jax._src.typing.SupportsDType | ~typing.Any = <class 'jax.numpy.float32'>, rngs: ~flax.nnx.rnglib.Rngs)