fix!: cond layer norm (#658)

Fix name of config param and doc string

<!-- readthedocs-preview anemoi-training start -->
----
📚 Documentation preview 📚:
https://anemoi-training--658.org.readthedocs.build/en/658/

<!-- readthedocs-preview anemoi-training end -->

<!-- readthedocs-preview anemoi-graphs start -->
----
📚 Documentation preview 📚:
https://anemoi-graphs--658.org.readthedocs.build/en/658/

<!-- readthedocs-preview anemoi-graphs end -->

<!-- readthedocs-preview anemoi-models start -->
----
📚 Documentation preview 📚:
https://anemoi-models--658.org.readthedocs.build/en/658/

<!-- readthedocs-preview anemoi-models end -->

---------

Co-authored-by: Ana Prieto Nemesio <[email protected]>

Author: ssmmnn11

models/src/anemoi/models/layers/normalization.py

@@ -42,16 +42,31 @@ def __init__(
         self,
         normalized_shape: Union[int, list, Size],
         condition_shape: int = 16,
-        w_one_bias_zero_init: bool = True,
+        zero_init: bool = True,
         autocast: bool = True,
-    ):
+    ) -> None:
+        """Initialize Conditional Layer Normalization.
+
+        Parameters
+        ----------
+        normalized_shape : Union[int, list, Size]
+            Shape or dimension(s) over which to normalize.
+        condition_shape : int, optional
+            Dimension of the conditioning vector, by default 16.
+        zero_init : bool, optional
+            If True, initializes the scale and bias transformation weights to zeros.
+            This means the conditional normalization behaves like standard layer
+            normalization initially, by default True.
+        autocast : bool, optional
+            If True, automatically cast output to match input dtype, by default True.
+        """
         super().__init__()
         self.norm = nn.LayerNorm(normalized_shape, elementwise_affine=False)  # no learnable parameters
         self.scale = nn.Linear(condition_shape, normalized_shape)  # , bias=False)
         self.bias = nn.Linear(condition_shape, normalized_shape)  # , bias=False)
         self.autocast = autocast
 
-        if w_one_bias_zero_init:
+        if zero_init:
             nn.init.zeros_(self.scale.weight)
             nn.init.zeros_(self.scale.bias)
             nn.init.zeros_(self.bias.weight)
@@ -62,15 +77,15 @@ def forward(self, x: Tensor, cond: Tensor) -> Tensor:
 
         Parameters
         ----------
-        input : List[Tensor, Tensor]
-            A list of two tensors (x, cond),
-            the first is the input tensor and
-            the second is the condition tensor.
+        x : Tensor
+            Input tensor to be normalized.
+        cond : Tensor
+            Conditioning tensor used to modulate the normalization.
 
         Returns
         -------
         Tensor
-            The output tensor.
+            Output tensor.
         """
         scale = self.scale(cond)
         bias = self.bias(cond)

training/docs/user-guide/diffusion-set-up.rst

@@ -87,7 +87,7 @@ The diffusion configuration includes:
        _target_: anemoi.models.layers.normalization.ConditionalLayerNorm
        normalized_shape: ${model.num_channels}
        condition_shape: 16
-       w_one_bias_zero_init: True
+       zero_init: True
        autocast: false
 
 The diffusion model uses conditional layer normalization to condition

training/docs/user-guide/kcrps-set-up.rst

@@ -113,7 +113,7 @@ conditional layer norm.
             _target_: anemoi.models.layers.normalization.ConditionalLayerNorm
             normalized_shape: ${model.num_channels}
             condition_shape: ${model.noise_injector.noise_channels_dim}
-            w_one_bias_zero_init: True
+            zero_init: True
             autocast: false
          ...
 

training/src/anemoi/training/config/model/graphtransformer_diffusion.yaml

@@ -37,7 +37,7 @@ layer_kernels:
     _target_: anemoi.models.layers.normalization.ConditionalLayerNorm
     normalized_shape: ${model.num_channels}
     condition_shape: 16
-    w_one_bias_zero_init: True
+    zero_init: True
     autocast: false
   Linear:
     _target_: torch.nn.Linear

training/src/anemoi/training/config/model/graphtransformer_diffusiontend.yaml

@@ -37,7 +37,7 @@ layer_kernels:
     _target_: anemoi.models.layers.normalization.ConditionalLayerNorm
     normalized_shape: ${model.num_channels}
     condition_shape: 16
-    w_one_bias_zero_init: True
+    zero_init: True
     autocast: false
   Linear:
     _target_: torch.nn.Linear

training/src/anemoi/training/config/model/graphtransformer_ens.yaml

@@ -48,7 +48,7 @@ processor:
       _target_: anemoi.models.layers.normalization.ConditionalLayerNorm
       normalized_shape: ${model.num_channels}
       condition_shape: ${model.noise_injector.noise_channels_dim}
-      w_one_bias_zero_init: True
+      zero_init: True
       autocast: false
       #Any arguments to your chosen function go here
     Linear:

training/src/anemoi/training/config/model/transformer_diffusion.yaml

@@ -37,7 +37,7 @@ layer_kernels:
     _target_: anemoi.models.layers.normalization.ConditionalLayerNorm
     normalized_shape: ${model.num_channels}
     condition_shape: 16
-    w_one_bias_zero_init: True
+    zero_init: True
     autocast: false
   Linear:
     _target_: torch.nn.Linear

training/src/anemoi/training/config/model/transformer_diffusiontend.yaml

@@ -37,7 +37,7 @@ layer_kernels:
     _target_: anemoi.models.layers.normalization.ConditionalLayerNorm
     normalized_shape: ${model.num_channels}
     condition_shape: 16
-    w_one_bias_zero_init: True
+    zero_init: True
     autocast: false
   Linear:
     _target_: torch.nn.Linear

training/src/anemoi/training/config/model/transformer_ens.yaml

@@ -50,7 +50,7 @@ processor:
       _target_: anemoi.models.layers.normalization.ConditionalLayerNorm
       normalized_shape: ${model.num_channels}
       condition_shape: ${model.noise_injector.noise_channels_dim}
-      w_one_bias_zero_init: True
+      zero_init: True
       autocast: false
       #Any arguments to your chosen function go here
     Linear: