From 40bb6b3b45e4657642f3d887961549224cde2750 Mon Sep 17 00:00:00 2001
From: PalaniappanR02 <palaniappan9c17@gmail.com>
Date: Sun, 22 Mar 2026 09:37:42 +0530
Subject: [PATCH 1/3] docs: standardize docstrings for ARModel init

---
 neural_lam/models/ar_model.py | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/neural_lam/models/ar_model.py b/neural_lam/models/ar_model.py
index f1bcb461d..f9291f718 100644
--- a/neural_lam/models/ar_model.py
+++ b/neural_lam/models/ar_model.py
@@ -36,6 +36,20 @@ def __init__(
         config: NeuralLAMConfig,
         datastore: BaseDatastore,
     ):
+        """
+        Initialize the Auto-Regressive model.
+
+        Parameters
+        ----------
+        args : argparse.Namespace
+            Command-line arguments containing model hyperparameters and 
+            run configurations.
+        config : NeuralLAMConfig
+            Configuration object for the NeuralLAM project.
+        datastore : BaseDatastore
+            Datastore object used to retrieve weather data and 
+            standardization statistics.
+        """
         super().__init__()
         self.save_hyperparameters(ignore=["datastore"])
         self.args = args

From 9e3986c02d782001d18a6f431c3c6f823288a6de Mon Sep 17 00:00:00 2001
From: PalaniappanR02 <palaniappan9c17@gmail.com>
Date: Sun, 22 Mar 2026 09:52:33 +0530
Subject: [PATCH 2/3] docs: standardize predict_step docstring in ARModel

---
 neural_lam/models/ar_model.py | 20 ++++++++++++++++----
 1 file changed, 16 insertions(+), 4 deletions(-)

diff --git a/neural_lam/models/ar_model.py b/neural_lam/models/ar_model.py
index f9291f718..a92760102 100644
--- a/neural_lam/models/ar_model.py
+++ b/neural_lam/models/ar_model.py
@@ -234,10 +234,22 @@ def expand_to_batch(x, batch_size):
 
     def predict_step(self, prev_state, prev_prev_state, forcing):
         """
-        Step state one step ahead using prediction model, X_{t-1}, X_t -> X_t+1
-        prev_state: (B, num_grid_nodes, feature_dim), X_t prev_prev_state: (B,
-        num_grid_nodes, feature_dim), X_{t-1} forcing: (B, num_grid_nodes,
-        forcing_dim)
+        Step state one step ahead using the prediction model.
+        Computes X_{t+1} given X_t and X_{t-1}.
+
+        Parameters
+        ----------
+        prev_state : torch.Tensor
+            State at time t (X_t), shape (B, num_grid_nodes, feature_dim).
+        prev_prev_state : torch.Tensor
+            State at time t-1 (X_{t-1}), shape (B, num_grid_nodes, feature_dim).
+        forcing : torch.Tensor
+            Forcing features at time t, shape (B, num_grid_nodes, forcing_dim).
+
+        Raises
+        ------
+        NotImplementedError
+            This is an abstract method that must be implemented by subclasses.
         """
         raise NotImplementedError("No prediction step implemented")
 

From 78b7dd0fb07b56f360071b74d7d8204351f610bc Mon Sep 17 00:00:00 2001
From: PalaniappanR02 <palaniappan9c17@gmail.com>
Date: Sun, 22 Mar 2026 10:04:45 +0530
Subject: [PATCH 3/3] docs: standardize mse docstring in metrics module

---
 neural_lam/metrics.py | 39 ++++++++++++++++++++++++---------------
 1 file changed, 24 insertions(+), 15 deletions(-)

diff --git a/neural_lam/metrics.py b/neural_lam/metrics.py
index 7db2cca6d..2adfdb542 100644
--- a/neural_lam/metrics.py
+++ b/neural_lam/metrics.py
@@ -86,21 +86,30 @@ def wmse(pred, target, pred_std, mask=None, average_grid=True, sum_vars=True):
 
 def mse(pred, target, pred_std, mask=None, average_grid=True, sum_vars=True):
     """
-    (Unweighted) Mean Squared Error
-
-    (...,) is any number of batch dimensions, potentially different
-        but broadcastable
-    pred: (..., N, d_state), prediction
-    target: (..., N, d_state), target
-    pred_std: (..., N, d_state) or (d_state,), predicted std.-dev.
-    mask: (N,), boolean mask describing which grid nodes to use in metric
-    average_grid: boolean, if grid dimension -2 should be reduced (mean over N)
-    sum_vars: boolean, if variable dimension -1 should be reduced (sum
-        over d_state)
-
-    Returns:
-    metric_val: One of (...,), (..., d_state), (..., N), (..., N, d_state),
-    depending on reduction arguments.
+    (Unweighted) Mean Squared Error.
+
+    Parameters
+    ----------
+    pred : torch.Tensor
+        Prediction tensor of shape (..., N, d_state).
+    target : torch.Tensor
+        Target tensor of shape (..., N, d_state).
+    pred_std : torch.Tensor
+        Predicted standard deviation of shape (..., N, d_state) or (d_state,).
+    mask : torch.Tensor, optional
+        Boolean mask of shape (N,) describing which grid nodes to use, 
+        by default None.
+    average_grid : bool, optional
+        Whether the grid dimension -2 should be reduced (mean over N), 
+        by default True.
+    sum_vars : bool, optional
+        Whether the variable dimension -1 should be reduced (sum over d_state), 
+        by default True.
+
+    Returns
+    -------
+    torch.Tensor
+        Metric value(s) depending on reduction arguments.
     """
     # Replace pred_std with constant ones
     return wmse(