attempt at mode check in one class

Author: ctr26

bioimageio/spec/model/v0_5.py

@@ -530,6 +530,22 @@ class ZeroMeanUnitVarianceKwargs(ProcessingKwargs):
     eps: Annotated[float, Interval(gt=0, le=0.1)] = 1e-6
     """epsilon for numeric stability: `out = (tensor - mean) / (std + eps)`."""
 
+    @model_validator(mode="after")
+    def validate_axes_based_on_mode(self)-> Self:
+        mode = self.mode
+        axes = self.axes
+
+        if mode == "per_sample":
+            # In 'per_sample' mode, ensure batch axis is not included in 'axes'
+            if axes and 'batch' in axes:
+                raise ValueError("Batch axis should not be included in 'axes' for 'per_sample' mode.")
+        elif mode == "per_dataset":
+            # In 'per_dataset' mode, check if batch axis handling is appropriate
+            # Implement any specific logic you need here
+            if axes and 'batch' not in axes:
+                raise ValueError("Batch axis must be included in 'axes' for 'per_dataset' mode.")
+        return self
+
 
 class ZeroMeanUnitVariance(ProcessingBase):
     """Subtract mean and divide by variance."""