docs: use Literal for strict typing of z-scoring options in Public API (#1744)

* refactor: use Literal for strict z-score typing

* fix the parser function

* revert internal modules to str to avoid type purging issues

* revert type hint to str to fix CI

Author: satwiksps

sbi/neural_nets/factory.py

@@ -76,8 +76,12 @@ class ZukoFlowType(Enum):
 
 def classifier_nn(
     model: str,
-    z_score_theta: Optional[str] = "independent",
-    z_score_x: Optional[str] = "independent",
+    z_score_theta: Optional[
+        Literal["independent", "structured", "transform_to_unconstrained", "none"]
+    ] = "independent",
+    z_score_x: Optional[
+        Literal["independent", "structured", "transform_to_unconstrained", "none"]
+    ] = "independent",
     hidden_features: int = 50,
     embedding_net_theta: nn.Module = nn.Identity(),
     embedding_net_x: nn.Module = nn.Identity(),
@@ -151,8 +155,12 @@ def build_fn(batch_theta, batch_x):
 
 def likelihood_nn(
     model: str,
-    z_score_theta: Optional[str] = "independent",
-    z_score_x: Optional[str] = "independent",
+    z_score_theta: Optional[
+        Literal["independent", "structured", "transform_to_unconstrained", "none"]
+    ] = "independent",
+    z_score_x: Optional[
+        Literal["independent", "structured", "transform_to_unconstrained", "none"]
+    ] = "independent",
     hidden_features: int = 50,
     num_transforms: int = 5,
     num_bins: int = 10,
@@ -226,8 +234,12 @@ def build_fn(batch_theta, batch_x):
 
 def posterior_nn(
     model: str,
-    z_score_theta: Optional[str] = "independent",
-    z_score_x: Optional[str] = "independent",
+    z_score_theta: Optional[
+        Literal["independent", "structured", "transform_to_unconstrained", "none"]
+    ] = "independent",
+    z_score_x: Optional[
+        Literal["independent", "structured", "transform_to_unconstrained", "none"]
+    ] = "independent",
     hidden_features: int = 50,
     num_transforms: int = 5,
     num_bins: int = 10,
@@ -334,8 +346,12 @@ def posterior_score_nn(
         VectorFieldNet,
     ] = "mlp",
     sde_type: str = "ve",
-    z_score_theta: Optional[str] = "independent",
-    z_score_x: Optional[str] = "independent",
+    z_score_theta: Optional[
+        Literal["independent", "structured", "transform_to_unconstrained", "none"]
+    ] = "independent",
+    z_score_x: Optional[
+        Literal["independent", "structured", "transform_to_unconstrained", "none"]
+    ] = "independent",
     hidden_features: int = 100,
     num_layers: int = 5,
     embedding_net: nn.Module = nn.Identity(),
@@ -436,8 +452,12 @@ def build_fn(batch_theta, batch_x):
 # TODO: remove this function on next release
 def flowmatching_nn(
     model: str,
-    z_score_theta: Optional[str] = "independent",
-    z_score_x: Optional[str] = "independent",
+    z_score_theta: Optional[
+        Literal["independent", "structured", "transform_to_unconstrained", "none"]
+    ] = "independent",
+    z_score_x: Optional[
+        Literal["independent", "structured", "transform_to_unconstrained", "none"]
+    ] = "independent",
     hidden_features: int = 64,
     num_layers: int = 5,
     num_blocks: int = 5,
@@ -510,8 +530,12 @@ def posterior_flow_nn(
         Literal["mlp", "ada_mlp", "transformer", "transformer_cross_attn"],
         VectorFieldNet,
     ] = "mlp",
-    z_score_theta: Optional[str] = None,
-    z_score_x: Optional[str] = "independent",
+    z_score_theta: Optional[
+        Literal["independent", "structured", "transform_to_unconstrained", "none"]
+    ] = None,
+    z_score_x: Optional[
+        Literal["independent", "structured", "transform_to_unconstrained", "none"]
+    ] = "independent",
     hidden_features: int = 100,
     num_layers: int = 5,
     embedding_net: nn.Module = nn.Identity(),
@@ -592,7 +616,9 @@ def build_fn(batch_theta, batch_x):
 
 def marginal_nn(
     model: ZukoFlowType,
-    z_score_x: Optional[str] = "independent",
+    z_score_x: Optional[
+        Literal["independent", "structured", "transform_to_unconstrained", "none"]
+    ] = "independent",
     hidden_features: int = 50,
     num_transforms: int = 5,
     num_bins: int = 10,

sbi/utils/restriction_estimator.py

@@ -3,7 +3,7 @@
 
 from copy import deepcopy
 from math import floor
-from typing import Any, Callable, Optional, Tuple, Union
+from typing import Any, Callable, Literal, Optional, Tuple, Union
 
 import torch
 import torch.nn.functional as F
@@ -74,7 +74,9 @@ def __init__(
         hidden_features: int = 100,
         num_blocks: int = 2,
         dropout_probability: float = 0.5,
-        z_score: Optional[str] = "independent",
+        z_score: Optional[
+            Literal["independent", "structured", "transform_to_unconstrained", "none"]
+        ] = "independent",
         embedding_net: nn.Module = nn.Identity(),
     ) -> None:
         r"""