diff --git a/pyproject.toml b/pyproject.toml
index 3c710eb..abf2491 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "scgenerator"
-version = "0.3.5"
+version = "0.3.6"
 description = "Simulate nonlinear pulse propagation in optical fibers"
 readme = "README.md"
 authors = [{ name = "Benoit Sierro", email = "benoit.sierro@iap.unibe.ch" }]
diff --git a/src/scgenerator/physics/pulse.py b/src/scgenerator/physics/pulse.py
index 4ece788..0a1733a 100644
--- a/src/scgenerator/physics/pulse.py
+++ b/src/scgenerator/physics/pulse.py
@@ -672,18 +672,19 @@ def g12(values):
 
     Parameters
     ----------
-    values : 2D array
+    values : np.ndarray, shape (..., m, n)
         complex values following sc-ordering
 
     Returns
     -------
-        g12_arr : coherence function as a n-D array
+    np.ndarray, shape (..., n)
+        coherence function
     """
 
     # Create all the possible pairs of values
     n = len(values)
     field_pairs = itertools.combinations(values, 2)
-    mean_spec = np.mean(math.abs2(values), axis=0)
+    mean_spec = np.mean(math.abs2(values), axis=-2)
     mask = mean_spec > 1e-15 * mean_spec.max()
     corr = np.zeros_like(values[0])
     for pair in field_pairs: