diff --git a/.mise.toml b/.mise.toml
index 04b145a77..8c21999c1 100644
--- a/.mise.toml
+++ b/.mise.toml
@@ -24,7 +24,7 @@ ripgrep = "latest"
 dprint = "latest"
 uv = "0.9.30"
 ruff = "0.15.0"
-ty = "0.0.32"
+ty = "0.0.44"
 
 "aqua:j178/prek" = "latest"
 
diff --git a/STYLE_GUIDE.md b/STYLE_GUIDE.md
index 8919bfeb6..199986840 100644
--- a/STYLE_GUIDE.md
+++ b/STYLE_GUIDE.md
@@ -141,6 +141,18 @@ def process(data: pd.DataFrame, columns: Sequence[str] | None = None) -> Self:
 - `Protocol` for structural subtyping when you need duck-typing boundaries
 - Avoid `Any` -- prefer `object`, generics, or `Protocol`
 - `TYPE_CHECKING` guards for heavy imports (`pandas`, `torch`, `transformers`); not needed for stdlib or lightweight imports
+- `TypeIs[T]` from `typing_extensions` for side-effect-free predicates that fully validate a value as `T`. Use `bool` for ordinary predicates that do not establish a type, and reserve `TypeGuard` for cases `TypeIs` cannot express because the narrowed type is not compatible with the input type.
+
+```python
+from typing_extensions import TypeIs
+
+
+def is_json_object(value: object) -> TypeIs[dict[str, JsonValue]]:
+    return isinstance(value, dict) and all(
+        isinstance(key, str) and is_json_value(item)
+        for key, item in value.items()
+    )
+```
 
 Legacy modules (`pii_replacer/`) still use `Optional`/`List`/`Dict`. Only `pii_replacer/` is excluded from `ty` type-checking (see `[tool.ty.src] exclude` in `pyproject.toml`). The remaining legacy usages will be migrated when those modules come under the type checker.
 
diff --git a/mise.lock b/mise.lock
index 04b34ec62..c83b7e0e3 100644
--- a/mise.lock
+++ b/mise.lock
@@ -131,22 +131,22 @@ checksum = "sha256:093d355ac33c6b8e91e80b8497d5581c61b028c0405e265cf38fd88f9a291
 url = "https://github.com/astral-sh/ruff/releases/download/0.15.0/ruff-aarch64-apple-darwin.tar.gz"
 
 [[tools.ty]]
-version = "0.0.32"
+version = "0.0.44"
 backend = "aqua:astral-sh/ty"
 
 [tools.ty."platforms.linux-arm64"]
-checksum = "sha256:de848acf867991f495dc346f5d12a7ac470af3ac464fdad7aff22fb8ee931a17"
-url = "https://github.com/astral-sh/ty/releases/download/0.0.32/ty-aarch64-unknown-linux-musl.tar.gz"
+checksum = "sha256:46a84ec41d2ae8892673f698a2ae31e02278203ededdce0d9229b7c4f5ac8bff"
+url = "https://github.com/astral-sh/ty/releases/download/0.0.44/ty-aarch64-unknown-linux-musl.tar.gz"
 provenance = "github-attestations"
 
 [tools.ty."platforms.linux-x64"]
-checksum = "sha256:cc58ee952aa551a0cbca495a43325b093dfcb180f0b314bf2c71068d37833b9e"
-url = "https://github.com/astral-sh/ty/releases/download/0.0.32/ty-x86_64-unknown-linux-musl.tar.gz"
+checksum = "sha256:618549cc5b0fd19b3ca0830a43fee4189623d7bd993bf65315f1f0ba650d944a"
+url = "https://github.com/astral-sh/ty/releases/download/0.0.44/ty-x86_64-unknown-linux-musl.tar.gz"
 provenance = "github-attestations"
 
 [tools.ty."platforms.macos-arm64"]
-checksum = "sha256:6b03b94d8c2ddcb5db67e6c863ef3c72b83fbcb973e0ff3a4c6daa4352dad009"
-url = "https://github.com/astral-sh/ty/releases/download/0.0.32/ty-aarch64-apple-darwin.tar.gz"
+checksum = "sha256:e796d5a91886a379d1da7c97772722a205901991d642382f75e2c27c138cfddb"
+url = "https://github.com/astral-sh/ty/releases/download/0.0.44/ty-aarch64-apple-darwin.tar.gz"
 provenance = "github-attestations"
 
 [[tools.uv]]
diff --git a/src/nemo_safe_synthesizer/cli/datasets.py b/src/nemo_safe_synthesizer/cli/datasets.py
index 7b8662aed..b206eb77c 100644
--- a/src/nemo_safe_synthesizer/cli/datasets.py
+++ b/src/nemo_safe_synthesizer/cli/datasets.py
@@ -21,6 +21,22 @@
 logger = get_logger(__name__)
 
 
+def _require_dataframe(value: object, *, url: str) -> pd.DataFrame:
+    if isinstance(value, pd.DataFrame):
+        return value
+    raise TypeError(f"Expected dataset reader for {url} to return a pandas DataFrame, got {type(value).__name__}")
+
+
+def _dynamic_callable(value: object) -> Any:
+    """Widen a typed pandas reader callable to ``Any``.
+
+    ty raises ``invalid-argument-type`` when a union of overloaded pandas reader
+    callables is dispatched with ``**kwargs``. Erasing the callable type here
+    keeps the runtime behavior unchanged while avoiding that diagnostic.
+    """
+    return value
+
+
 class DatasetInfo(BaseModel):
     """Entry in the dataset registry."""
 
@@ -96,20 +112,22 @@ def fetch(self) -> pd.DataFrame:
         logger.info(f"Reading dataset from {url}")
 
         # Determine the file extension and appropriate reader
-        match Path(url).suffix.lstrip("."):
+        reader: Any
+        extension = Path(url).suffix.lstrip(".")
+        match extension:
             case "csv" | "txt":
-                reader = pd.read_csv
+                reader = _dynamic_callable(pd.read_csv)
                 default_load_args: dict[str, Any] = {}
             case "json":
-                reader = pd.read_json
+                reader = _dynamic_callable(pd.read_json)
                 default_load_args = {}
             case "jsonl":
-                reader = pd.read_json
+                reader = _dynamic_callable(pd.read_json)
                 default_load_args = {"lines": True}
             case "parquet":
-                reader = pd.read_parquet
+                reader = _dynamic_callable(pd.read_parquet)
                 default_load_args = {}
-            case extension:
+            case _:
                 if not extension:
                     extension = f"<no extension found on url '{url}'>"
                 raise ValueError(f"Unsupported file extension: {extension}")
@@ -118,7 +136,7 @@ def fetch(self) -> pd.DataFrame:
         final_load_args = {**default_load_args, **(self.load_args or {})}
 
         try:
-            return reader(url, **final_load_args)  # ty: ignore[invalid-argument-type] -- reader union includes parquet which has stricter signature
+            return _require_dataframe(reader(url, **final_load_args), url=url)
         except Exception as e:
             logger.error(f"Error reading dataset from {url}: {e}", exc_info=True)
             raise
diff --git a/src/nemo_safe_synthesizer/data_processing/actions/dates.py b/src/nemo_safe_synthesizer/data_processing/actions/dates.py
index 3575391d4..99ac20c25 100644
--- a/src/nemo_safe_synthesizer/data_processing/actions/dates.py
+++ b/src/nemo_safe_synthesizer/data_processing/actions/dates.py
@@ -13,7 +13,7 @@
 import itertools
 import re
 from collections import Counter
-from collections.abc import Iterable, Iterator
+from collections.abc import Hashable, Iterable, Iterator
 from dataclasses import dataclass
 from datetime import datetime, timedelta
 from random import randint
@@ -399,7 +399,7 @@ def infer_from_series(date_series: Iterable[str]) -> Optional[str]:
 def fit_and_transform_dates(
     df: pd.DataFrame,
     inplace: bool = False,
-) -> tuple[dict[str, dict[str, str]], pd.DataFrame]:
+) -> tuple[dict[Hashable, dict[str, str]], pd.DataFrame]:
     """Detect date columns, convert them to elapsed seconds, and record the transformation.
 
     For each object-typed column, samples values to infer a date format. If
@@ -412,10 +412,10 @@ def fit_and_transform_dates(
 
     Returns:
         A tuple of (date_min_dict, result_df). ``date_min_dict`` maps column
-        names to ``{"format": ..., "min": ...}`` dicts needed by
+        labels to ``{"format": ..., "min": ...}`` dicts needed by
         ``transform_dates`` for reversal.
     """
-    date_min_dict = {}
+    date_min_dict: dict[Hashable, dict[str, str]] = {}
     object_cols = [col for col, col_type in df.dtypes.items() if col_type == "object"]
     result_df = df.copy() if not inplace else df
     for object_col in object_cols:
@@ -437,11 +437,11 @@ def fit_and_transform_dates(
     return date_min_dict, result_df
 
 
-def transform_dates(dates: dict[str, dict[str, str]], df: pd.DataFrame) -> pd.DataFrame:
+def transform_dates(dates: dict[Hashable, dict[str, str]], df: pd.DataFrame) -> pd.DataFrame:
     """Apply a previously fitted date-to-seconds transformation to a DataFrame.
 
     Args:
-        dates: Mapping from column names to ``{"format": ..., "min": ...}``
+        dates: Mapping from column labels to ``{"format": ..., "min": ...}``
             dicts as returned by ``fit_and_transform_dates``.
         df: DataFrame to transform.
 
diff --git a/src/nemo_safe_synthesizer/data_processing/actions/distributions.py b/src/nemo_safe_synthesizer/data_processing/actions/distributions.py
index 2d9c6e499..7bb325dd8 100644
--- a/src/nemo_safe_synthesizer/data_processing/actions/distributions.py
+++ b/src/nemo_safe_synthesizer/data_processing/actions/distributions.py
@@ -13,7 +13,6 @@
 
 from abc import ABC, abstractmethod
 from datetime import datetime, timedelta
-from functools import partial
 from typing import Annotated, Any, Literal, Optional, Union
 
 import numpy as np
@@ -77,21 +76,20 @@ def _round_datetime(self, dt: datetime, precision: timedelta) -> datetime:
         rounded_ts = round(dt.timestamp() / precision.total_seconds()) * precision.total_seconds()
         return datetime.fromtimestamp(rounded_ts)
 
-    def _apply_universal_params(self, samples: list[datetime]) -> list[datetime]:
-        ret: list[str] | list[datetime] = samples
-
-        ops = []
-        if self.precision is not None:
-            ops.append(partial(self._round_datetime, self.precision))
+    def _apply_universal_params(self, samples: list[datetime]) -> list[datetime] | list[str]:
         if self.format is not None:
-            ops.append(lambda x: x.strftime(self.format))
-
+            formatted: list[str] = []
+            for sample in samples:
+                if self.precision is not None:
+                    sample = self._round_datetime(sample, self.precision)
+                formatted.append(sample.strftime(self.format))
+            return formatted
+
+        ret: list[datetime] = []
         for sample in samples:
-            n = sample
-            for op in ops:
-                n = op(n)
-            ret.append(n)
-
+            if self.precision is not None:
+                sample = self._round_datetime(sample, self.precision)
+            ret.append(sample)
         return ret
 
 
diff --git a/src/nemo_safe_synthesizer/evaluation/components/attribute_inference_protection.py b/src/nemo_safe_synthesizer/evaluation/components/attribute_inference_protection.py
index 2a9c6e4bd..0fcc9e246 100644
--- a/src/nemo_safe_synthesizer/evaluation/components/attribute_inference_protection.py
+++ b/src/nemo_safe_synthesizer/evaluation/components/attribute_inference_protection.py
@@ -6,6 +6,7 @@
 import itertools
 import math
 import re
+from collections.abc import Hashable
 from datetime import datetime
 from decimal import Decimal
 from functools import cached_property
@@ -138,14 +139,15 @@ def _pandas_entropy(column: pd.Series, base: float | None = None) -> np.float64:
         return -(vc * np.log(vc) / np.log(base)).sum()
 
     @staticmethod
-    def _is_really_categorical(column: str) -> bool:
+    def _is_really_categorical(column: Hashable) -> bool:
+        column_name = str(column)
         # Break the header up into parts
         for separator in ["_", " ", "-", "."]:
-            if separator in column:
-                col_parts = column.split(separator)
+            if separator in column_name:
+                col_parts = column_name.split(separator)
                 break
             else:
-                col_parts = [column]
+                col_parts = [column_name]
 
         # Go through the parts and divide up camel case
         col_final_parts = []
@@ -366,16 +368,16 @@ def _aia(
 
             # Get all combinations of columns to be the quasi-identifiers
             # This gets explosive when column count > 500
+            training_columns = list(training_df.columns)
             if len(training_df.columns) < 500:
-                qi_combos = list(itertools.combinations(training_df.columns, quasi_identifier_count))
+                qi_combos = list(itertools.combinations(training_columns, quasi_identifier_count))
             else:
-                columns = list(training_df.columns)
                 qi_combos = []
-                for i in range(len(columns) - quasi_identifier_count):
-                    combo = set()
+                for i in range(len(training_columns) - quasi_identifier_count + 1):
+                    combo = []
                     for j in range(quasi_identifier_count):
-                        combo.add(columns[i + j])
-                    qi_combos.append(combo)
+                        combo.append(training_columns[i + j])
+                    qi_combos.append(tuple(combo))
 
             np.random.seed(5)
             np.random.shuffle(qi_combos)
@@ -415,7 +417,6 @@ def _aia(
 
             # As we process the attack dataset, we'll accumulate for each column the number of
             # correct and incorrect predictions
-            training_columns = [str(column) for column in training_df.columns]
             correct = {predict_column: 0 for predict_column in training_columns}
             incorrect = {predict_column: 0 for predict_column in training_columns}
 
@@ -431,14 +432,15 @@ def _aia(
                     more_to_process = False
                     continue
 
-                # Randomly sample columns to be the `qi`
-                qi = qi_combos[qi_index]
-                qi_index += 1
                 # We stop processing if all qi combos have been processed.
                 if qi_index == len(qi_combos):
                     more_to_process = False
                     continue
 
+                # Randomly sample columns to be the `qi`
+                qi = qi_combos[qi_index]
+                qi_index += 1
+
                 # Predict columns are all but the `qi`
                 predict_columns = [column for column in training_columns if column not in qi]
 
@@ -578,7 +580,8 @@ def _aia(
                 for i in range(len(entropy)):
                     entropy_wts.append(0)
             else:
-                arr = (entropy - min(entropy)) / (max(entropy) - min(entropy))
+                entropy_arr = np.asarray(entropy, dtype=float)
+                arr = (entropy_arr - min(entropy)) / (max(entropy) - min(entropy))
                 entropy_wts = arr / arr.sum()
 
             i = 0
diff --git a/src/nemo_safe_synthesizer/evaluation/components/membership_inference_protection.py b/src/nemo_safe_synthesizer/evaluation/components/membership_inference_protection.py
index b307615f3..1490eb0e8 100644
--- a/src/nemo_safe_synthesizer/evaluation/components/membership_inference_protection.py
+++ b/src/nemo_safe_synthesizer/evaluation/components/membership_inference_protection.py
@@ -248,8 +248,8 @@ def _compute_mia(
     ) -> tuple[
         float,
         list[str],
-        dict[str, list[int]],
-        dict[str, list[int]],
+        dict[float, int],
+        dict[float, int],
     ]:
         """Core membership inference attack implementation for a single run.
 
@@ -286,8 +286,8 @@ def _compute_mia(
             pd.concat([training_df_attack, test_df_norm]).reset_index(drop=True).sample(frac=1, random_state=run)
         )
 
-        attack_synth_dist_text = [[0] for i in range(len(attack_df))]
-        attack_synth_indices_text = [[0] for i in range(len(attack_df))]
+        attack_synth_dist_text: list[list[float]] = [[0.0] for i in range(len(attack_df))]
+        attack_synth_indices_text: list[list[int]] = [[0] for i in range(len(attack_df))]
 
         # Get the NN dist for text for the entire attack dataset
 
@@ -362,8 +362,8 @@ def _compute_mia(
         score = 0
         attack_summary = []
 
-        tp_cnts = {}
-        fp_cnts = {}
+        tp_cnts: dict[float, int] = {}
+        fp_cnts: dict[float, int] = {}
 
         # Using the above text and tabular distances we now compute an overall distance score for
         # every record in the attack dataset. We then conduct 36 individual mia attacks on this one big
@@ -519,8 +519,8 @@ def mia(
 
             scores = []
             attack_sum_values = []
-            tps_values = {}
-            fps_values = {}
+            tps_values: dict[float, int] = {}
+            fps_values: dict[float, int] = {}
             for i in [0.1, 0.2, 0.3, 0.4]:
                 tps_values[i] = 0
                 fps_values[i] = 0
diff --git a/src/nemo_safe_synthesizer/evaluation/components/multi_modal_figures.py b/src/nemo_safe_synthesizer/evaluation/components/multi_modal_figures.py
index 4b3ed1fc9..d04f2ce95 100644
--- a/src/nemo_safe_synthesizer/evaluation/components/multi_modal_figures.py
+++ b/src/nemo_safe_synthesizer/evaluation/components/multi_modal_figures.py
@@ -676,7 +676,7 @@ def generate_text_structure_similarity_figures(
         "average_words_per_sentence",
         "average_characters_per_word",
     ]
-    figures = []
+    figures: list[go.Figure] = []
     for key in statistics_keys:
         if training_statistics.per_record_statistics.empty or synthetic_statistics.per_record_statistics.empty:
             break
@@ -684,7 +684,8 @@ def generate_text_structure_similarity_figures(
             training_statistics.per_record_statistics[key],
             synthetic_statistics.per_record_statistics[key],
         )
-        figures.append(figure)
+        if figure is not None:
+            figures.append(figure)
     if not figures:
         return None
 
diff --git a/src/nemo_safe_synthesizer/evaluation/reports/multimodal/multimodal_report.py b/src/nemo_safe_synthesizer/evaluation/reports/multimodal/multimodal_report.py
index b0d7358eb..90410d1b7 100644
--- a/src/nemo_safe_synthesizer/evaluation/reports/multimodal/multimodal_report.py
+++ b/src/nemo_safe_synthesizer/evaluation/reports/multimodal/multimodal_report.py
@@ -21,6 +21,7 @@
     ColumnDistribution,
     ColumnDistributionPlotRow,
 )
+from ....evaluation.components.component import Component
 from ....evaluation.components.correlation import (
     Correlation,
 )
@@ -150,7 +151,7 @@ def from_dataframes(
             mandatory_columns=MultimodalReport._get_config_value("mandatory_columns", [], config),
         )
 
-        components = []
+        components: list[Component] = []
 
         attribute_inference_protection = AttributeInferenceProtection(
             score=EvaluationScore(grade=PrivacyGrade.UNAVAILABLE)
diff --git a/src/nemo_safe_synthesizer/privacy/dp_transformers/linear.py b/src/nemo_safe_synthesizer/privacy/dp_transformers/linear.py
index 3c5473412..2f566030c 100644
--- a/src/nemo_safe_synthesizer/privacy/dp_transformers/linear.py
+++ b/src/nemo_safe_synthesizer/privacy/dp_transformers/linear.py
@@ -22,6 +22,19 @@
 from opt_einsum import contract
 
 
+def _contract_tensor(expression: str, *operands: torch.Tensor) -> torch.Tensor:
+    result = contract(expression, *operands)
+    if not isinstance(result, torch.Tensor):
+        raise TypeError("expected opt_einsum.contract to return a torch.Tensor for torch operands")
+    return result
+
+
+def _linear_parameter(parameter: torch.Tensor) -> nn.Parameter:
+    if not isinstance(parameter, nn.Parameter):
+        raise TypeError("expected nn.Linear parameter")
+    return parameter
+
+
 @register_grad_sampler(nn.Linear)
 def compute_linear_grad_sample(
     layer: nn.Linear, activations: list[torch.Tensor], backprops: torch.Tensor
@@ -41,10 +54,12 @@ def compute_linear_grad_sample(
         per-sample gradient tensor of shape ``(batch, ...)``.
     """
     activation = activations[0]
-    ret = {}
-    if layer.weight.requires_grad:
-        gs = contract("n...i,n...j->nij", backprops.float(), activation.float())
-        ret[layer.weight] = gs
-    if layer.bias is not None and layer.bias.requires_grad:
-        ret[layer.bias] = contract("n...k->nk", backprops.float())
+    ret: dict[nn.Parameter, torch.Tensor] = {}
+    weight = _linear_parameter(layer.weight)
+    if weight.requires_grad:
+        ret[weight] = _contract_tensor("n...i,n...j->nij", backprops.float(), activation.float())
+    if layer.bias is not None:
+        bias = _linear_parameter(layer.bias)
+        if bias.requires_grad:
+            ret[bias] = _contract_tensor("n...k->nk", backprops.float())
     return ret
diff --git a/src/nemo_safe_synthesizer/telemetry.py b/src/nemo_safe_synthesizer/telemetry.py
index 2702c75e7..5602b98ca 100644
--- a/src/nemo_safe_synthesizer/telemetry.py
+++ b/src/nemo_safe_synthesizer/telemetry.py
@@ -23,7 +23,7 @@
 from datetime import datetime, timezone
 from enum import Enum
 from pathlib import Path, PureWindowsPath
-from typing import TYPE_CHECKING, Any, ClassVar, cast
+from typing import TYPE_CHECKING, Any, ClassVar
 from urllib.parse import urlsplit, urlunsplit
 
 from pydantic import BaseModel, Field
@@ -78,7 +78,7 @@ def _redact_endpoint(endpoint: str) -> str:
     except ValueError:
         return "<invalid-endpoint>"
     query = "<redacted>" if parsed.query else ""
-    return cast(str, urlunsplit((parsed.scheme, parsed.netloc, parsed.path, query, parsed.fragment)))
+    return urlunsplit((parsed.scheme, parsed.netloc, parsed.path, query, parsed.fragment))
 
 
 def _deployment_type() -> DeploymentTypeEnum:
@@ -169,7 +169,7 @@ class NSSTrainingAndGenerationEvent(BaseModel):
 
     nemo_source: NemoSourceEnum = Field(
         default=NemoSourceEnum.SAFE_SYNTHESIZER,
-        alias="nemoSource",
+        serialization_alias="nemoSource",
         description="The NeMo product that created the event.",
     )
     task: str = Field(
@@ -183,72 +183,72 @@ class NSSTrainingAndGenerationEvent(BaseModel):
     )
     deployment_type: DeploymentTypeEnum = Field(
         default_factory=_deployment_type,
-        alias="deploymentType",
+        serialization_alias="deploymentType",
         description="How Safe Synthesizer was invoked (cli, sdk, nmp).",
     )
 
     # Timing
     job_duration_sec: float = Field(
         default=-1.0,
-        alias="jobDurationSec",
+        serialization_alias="jobDurationSec",
         description="Wall-clock duration of the job in seconds. -1.0 if not available.",
     )
 
     # Generation metrics
     num_records_generated: int = Field(
         default=-1,
-        alias="numRecordsGenerated",
+        serialization_alias="numRecordsGenerated",
         description="Number of valid synthetic records produced. -1 if not available.",
     )
     num_tokens_generated: int = Field(
         default=-1,
-        alias="numTokensGenerated",
+        serialization_alias="numTokensGenerated",
         description="Number of tokens generated by the model. -1 if not available.",
     )
 
     # Feature flags
     replace_pii_enabled: bool = Field(
         default=False,
-        alias="replacePiiEnabled",
+        serialization_alias="replacePiiEnabled",
         description="Whether PII replacement was enabled for this run.",
     )
     differential_privacy_enabled: bool = Field(
         default=False,
-        alias="differentialPrivacyEnabled",
+        serialization_alias="differentialPrivacyEnabled",
         description="Whether differential privacy training was enabled for this run.",
     )
     time_series_enabled: bool = Field(
         default=False,
-        alias="timeSeriesEnabled",
+        serialization_alias="timeSeriesEnabled",
         description="Whether time-series mode was enabled for this run.",
     )
     group_by_enabled: bool = Field(
         default=False,
-        alias="groupByEnabled",
+        serialization_alias="groupByEnabled",
         description="Whether group-by was set on the input data for this run.",
     )
 
     # Input characteristics (bucketed to avoid transmitting exact counts)
     input_records_bucket: str = Field(
         default="undefined",
-        alias="inputRecordsBucket",
+        serialization_alias="inputRecordsBucket",
         description="Bucketed count of input training records (e.g. '101-1000'). Use bucket_records().",
     )
     input_columns_bucket: str = Field(
         default="undefined",
-        alias="inputColumnsBucket",
+        serialization_alias="inputColumnsBucket",
         description="Bucketed count of input columns (e.g. '6-10'). Use bucket_columns().",
     )
 
     # Evaluation scores (-1.0 when evaluation was skipped or unavailable)
     synthetic_quality_score: float = Field(
         default=-1.0,
-        alias="syntheticQualityScore",
+        serialization_alias="syntheticQualityScore",
         description="Top-level Synthetic Quality Score from the evaluation report. -1.0 if not available.",
     )
     data_privacy_score: float = Field(
         default=-1.0,
-        alias="dataPrivacyScore",
+        serialization_alias="dataPrivacyScore",
         description="Top-level Data Privacy Score from the evaluation report. -1.0 if not available.",
     )
 
@@ -262,8 +262,6 @@ class NSSTrainingAndGenerationEvent(BaseModel):
         description="GPU device name (e.g. 'NVIDIA A100 80GB PCIe'). 'undefined' if not on GPU.",
     )
 
-    model_config = {"populate_by_name": True}
-
 
 @dataclass
 class QueuedEvent:
diff --git a/src/nemo_safe_synthesizer/utils.py b/src/nemo_safe_synthesizer/utils.py
index 93ac5bac8..49413973f 100644
--- a/src/nemo_safe_synthesizer/utils.py
+++ b/src/nemo_safe_synthesizer/utils.py
@@ -20,6 +20,7 @@
 import numpy as np
 import pandas as pd
 from pandas import DataFrame
+from typing_extensions import TypeIs
 
 from .data_processing.stats import Statistics
 from .observability import get_logger
@@ -37,6 +38,10 @@
 _HF_OFFLINE_ENV_VARS = ("HF_HUB_OFFLINE", "TRANSFORMERS_OFFLINE")
 
 
+def _is_statistics_list(stats: Statistics | list[Statistics]) -> TypeIs[list[Statistics]]:
+    return isinstance(stats, list) and all(isinstance(stat, Statistics) for stat in stats)
+
+
 def env_flag_is_true(name: str, *, default: bool = False) -> bool:
     """Return whether ``name`` is set to a truthy env value.
 
@@ -125,11 +130,11 @@ def log_stats(
         title: Optional table title.
     """
     headers = headers or []
-    stats = stats if isinstance(stats, list) else [stats]
+    stats_list = stats if _is_statistics_list(stats) else [stats]
 
     # Build structured data - processor will render as table for console
     structured_stats = {}
-    for header, stat in zip(headers, stats):
+    for header, stat in zip(headers, stats_list):
         key = header.lower().replace(" ", "_")
         structured_stats[key] = {
             "min": round_number_if_float(stat.min),
diff --git a/tests/cli/test_datasets.py b/tests/cli/test_datasets.py
index d85552ea4..9246e1744 100644
--- a/tests/cli/test_datasets.py
+++ b/tests/cli/test_datasets.py
@@ -181,6 +181,21 @@ def test_fetch_with_custom_load_args(self, tmp_path: Path):
         assert list(result.columns) == ["col1", "col2"]
         assert len(result) == 2
 
+    def test_fetch_rejects_non_dataframe_reader_result(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
+        """Test fetch raises a clear error when a pandas reader returns an unexpected type."""
+        csv_file = tmp_path / "test.csv"
+        csv_file.write_text("col1,col2\n1,2\n")
+
+        def fake_read_csv(*_args, **_kwargs) -> list[dict[str, int]]:
+            return [{"col1": 1, "col2": 2}]
+
+        monkeypatch.setattr(pd, "read_csv", fake_read_csv)
+
+        info = DatasetInfo(name="test", url=str(csv_file))
+
+        with pytest.raises(TypeError, match="Expected dataset reader"):
+            info.fetch()
+
     def test_fetch_uses_registry_base_url(self, tmp_path: Path):
         """Test fetch uses get_url() which includes base_url."""
         data_dir = tmp_path / "data"
diff --git a/tests/config/test_autoconfig.py b/tests/config/test_autoconfig.py
index 189b2627f..aa9c38508 100644
--- a/tests/config/test_autoconfig.py
+++ b/tests/config/test_autoconfig.py
@@ -72,9 +72,9 @@ class AutoConfigTestCase:
 
     def get_config(self) -> SafeSynthesizerParameters:
         """Get the config, calling it if it's a factory function."""
-        if callable(self.config):
-            return self.config()  # ty: ignore[call-top-callable] -- dynamic callable
-        return self.config
+        if isinstance(self.config, SafeSynthesizerParameters):
+            return self.config
+        return self.config()
 
 
 AUTO_NO_DP = AutoConfigTestCase(
@@ -170,6 +170,13 @@ def get_config(self) -> SafeSynthesizerParameters:
 ]
 
 
+def test_auto_config_test_case_get_config_calls_factory():
+    config = SafeSynthesizerParameters()
+    test_case = AutoConfigTestCase(name="factory", config=lambda: config, expected=AUTO_NO_DP.expected)
+
+    assert test_case.get_config() is config
+
+
 @pytest.fixture
 def sample_data() -> pd.DataFrame:
     """Standard test DataFrame (100 rows)."""
diff --git a/tests/data_processing/test_dates.py b/tests/data_processing/test_dates.py
new file mode 100644
index 000000000..657589066
--- /dev/null
+++ b/tests/data_processing/test_dates.py
@@ -0,0 +1,40 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import pandas as pd
+
+from nemo_safe_synthesizer.data_processing.actions.dates import fit_and_transform_dates, transform_dates
+
+
+def test_fit_and_transform_dates_records_detected_date_column():
+    df = pd.DataFrame(
+        {
+            "event_date": ["2024-01-01", "2024-01-03"],
+            "label": ["start", "end"],
+        }
+    )
+
+    date_columns, transformed = fit_and_transform_dates(df)
+
+    assert date_columns["event_date"]["format"] == "%Y-%m-%d"
+    assert date_columns["event_date"]["min"].startswith("2024-01-01")
+    assert transformed["event_date"].tolist() == [0.0, 172800.0]
+    assert df["event_date"].tolist() == ["2024-01-01", "2024-01-03"]
+
+
+def test_fit_and_transform_dates_preserves_non_string_column_label():
+    df = pd.DataFrame(
+        {
+            0: ["2024-01-01", "2024-01-03"],
+            "label": ["start", "end"],
+        }
+    )
+
+    date_columns, transformed = fit_and_transform_dates(df)
+
+    assert 0 in date_columns
+    assert "0" not in date_columns
+    assert transformed[0].tolist() == [0.0, 172800.0]
+    assert transform_dates(date_columns, df)[0].tolist() == [0.0, 172800.0]
diff --git a/tests/data_processing/test_distributions.py b/tests/data_processing/test_distributions.py
new file mode 100644
index 000000000..e6eedcacb
--- /dev/null
+++ b/tests/data_processing/test_distributions.py
@@ -0,0 +1,27 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from datetime import datetime, timedelta
+
+from nemo_safe_synthesizer.data_processing.actions.distributions import DatetimeDistribution
+
+
+class FixedDatetimeDistribution(DatetimeDistribution):
+    def sample_datetimes(self, num_records: int) -> list[datetime]:
+        start = datetime(2024, 1, 1, 12, 20)
+        return [start + timedelta(minutes=20 * offset) for offset in range(num_records)]
+
+
+def test_datetime_distribution_applies_precision_and_format_without_mutating_samples():
+    distribution = FixedDatetimeDistribution(precision=timedelta(hours=1), format="%H:%M")
+
+    assert distribution.sample(2) == ["12:00", "13:00"]
+
+
+def test_datetime_distribution_applies_precision_without_format():
+    distribution = FixedDatetimeDistribution(precision=timedelta(hours=1))
+
+    assert distribution.sample(2) == [
+        datetime(2024, 1, 1, 12),
+        datetime(2024, 1, 1, 13),
+    ]
diff --git a/tests/evaluation/components/test_attribute_inference_protection.py b/tests/evaluation/components/test_attribute_inference_protection.py
index a1771766a..9be5b776c 100644
--- a/tests/evaluation/components/test_attribute_inference_protection.py
+++ b/tests/evaluation/components/test_attribute_inference_protection.py
@@ -4,6 +4,8 @@
 # ruff: noqa: E402
 import logging
 
+import numpy as np
+import pandas as pd
 import pytest
 
 # Skip all tests in this module if sentence_transformers is not available
@@ -18,6 +20,91 @@
 logger = logging.getLogger(__name__)
 
 
+def test_aia_tabular_unit_exercises_entropy_weighting(monkeypatch: pytest.MonkeyPatch):
+    """Cover the fast tabular AIA path without loading text embedding models."""
+    training_df = pd.DataFrame(
+        {
+            "stable": [1, 1, 1, 1],
+            "binary": [1, 1, 2, 2],
+            "varied": [1, 2, 3, 4],
+        }
+    )
+    synthetic_df = training_df.copy()
+
+    def fake_get_synth_nn(*_args, **_kwargs) -> pd.DataFrame:
+        return synthetic_df.head(2)
+
+    monkeypatch.setattr(AttributeInferenceProtection, "_get_synth_nn", staticmethod(fake_get_synth_nn))
+
+    score, col_accuracy_df = AttributeInferenceProtection._aia(
+        training_df=training_df,
+        synthetic_df=synthetic_df,
+        quasi_identifier_count=1,
+    )
+
+    assert score.score is not None
+    assert col_accuracy_df is not None
+    assert list(col_accuracy_df["Column"]) == ["stable", "binary", "varied"]
+
+
+def test_aia_wide_table_uses_windowed_quasi_identifier_combinations(monkeypatch: pytest.MonkeyPatch):
+    """Cover the wide-table fallback that avoids materializing all combinations."""
+    columns = [f"col_{index}" for index in range(501)]
+    values = np.vstack([np.arange(501), np.arange(501) + 1])
+    training_df = pd.DataFrame(values, columns=columns)
+    synthetic_df = training_df.copy()
+    synth_calls: list[tuple[str, ...]] = []
+
+    def fake_get_synth_nn(train_row, *_args, **_kwargs) -> pd.DataFrame:
+        synth_calls.append(tuple(train_row.columns))
+        return synthetic_df.head(2)
+
+    monkeypatch.setattr(AttributeInferenceProtection, "_get_synth_nn", staticmethod(fake_get_synth_nn))
+
+    score, col_accuracy_df = AttributeInferenceProtection._aia(
+        training_df=training_df,
+        synthetic_df=synthetic_df,
+        quasi_identifier_count=500,
+    )
+
+    assert score.score is not None
+    assert col_accuracy_df is not None
+    assert len(col_accuracy_df) == len(columns)
+    assert len(synth_calls) == 2
+    assert set(synth_calls) == {tuple(columns[:500]), tuple(columns[1:])}
+
+
+def test_aia_preserves_non_string_column_labels(monkeypatch: pytest.MonkeyPatch):
+    """Cover DataFrames whose column labels are not strings."""
+    training_df = pd.DataFrame(
+        {
+            0: [1, 2, 3, 4],
+            1: [10, 20, 30, 40],
+            2: [100, 200, 300, 400],
+        }
+    )
+    synthetic_df = training_df.copy()
+    synth_calls: list[tuple[int, ...]] = []
+
+    def fake_get_synth_nn(train_row, *_args, **_kwargs) -> pd.DataFrame:
+        synth_calls.append(tuple(train_row.columns))
+        return synthetic_df.head(2)
+
+    monkeypatch.setattr(AttributeInferenceProtection, "_get_synth_nn", staticmethod(fake_get_synth_nn))
+
+    score, col_accuracy_df = AttributeInferenceProtection._aia(
+        training_df=training_df,
+        synthetic_df=synthetic_df,
+        quasi_identifier_count=1,
+    )
+
+    assert score.score is not None
+    assert col_accuracy_df is not None
+    assert set(col_accuracy_df["Column"]) == {0, 1, 2}
+    assert synth_calls
+    assert all(isinstance(column, int) for call in synth_calls for column in call)
+
+
 @pytest.mark.slow
 def test_attribute_inference_protection(fixture_training_df_5k, fixture_synthetic_df_5k, fixture_test_df):
     """Test AIA with tabular-only data (sklearn NearestNeighbors path)."""
diff --git a/tests/evaluation/components/test_membership_inference_protection.py b/tests/evaluation/components/test_membership_inference_protection.py
index c40983ad4..e26f82514 100644
--- a/tests/evaluation/components/test_membership_inference_protection.py
+++ b/tests/evaluation/components/test_membership_inference_protection.py
@@ -1,6 +1,8 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+import numpy as np
+import pandas as pd
 import pytest
 
 # Skip all tests in this module if sentence_transformers is not available
@@ -17,6 +19,34 @@
 logger = logging.getLogger(__name__)
 
 
+def test_mia_tabular_unit_returns_threshold_counts():
+    """Cover the fast tabular MIA path without loading text embedding models."""
+    training_df = pd.DataFrame(
+        {
+            "x": np.arange(24, dtype=float),
+            "y": np.arange(24, dtype=float) * 2,
+        }
+    )
+    synthetic_df = training_df.head(6).reset_index(drop=True)
+    test_df = pd.DataFrame(
+        {
+            "x": np.arange(100, 112, dtype=float),
+            "y": np.arange(100, 112, dtype=float) * 2,
+        }
+    )
+
+    score, attack_sum_df, tps_values, fps_values = MembershipInferenceProtection.mia(
+        training_df=training_df,
+        synthetic_df=synthetic_df,
+        test_df=test_df,
+    )
+
+    assert score.score is not None
+    assert attack_sum_df is not None
+    assert set(tps_values) == {0.1, 0.2, 0.3, 0.4}
+    assert set(fps_values) == {0.1, 0.2, 0.3, 0.4}
+
+
 @pytest.mark.requires_gpu
 def test_membership_inference_protection(fixture_training_df_5k, fixture_synthetic_df_5k, fixture_test_df):
     """Test MIA with tabular-only data (sklearn NearestNeighbors path)."""
diff --git a/tests/test_utils.py b/tests/test_utils.py
new file mode 100644
index 000000000..de9ef1d44
--- /dev/null
+++ b/tests/test_utils.py
@@ -0,0 +1,15 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from typing import Any, cast
+
+from nemo_safe_synthesizer.data_processing.stats import Statistics
+from nemo_safe_synthesizer.utils import _is_statistics_list
+
+
+def test_is_statistics_list_validates_all_items():
+    assert _is_statistics_list([Statistics()])
+    assert not _is_statistics_list(Statistics())
+    assert not _is_statistics_list(cast(Any, [object()]))
diff --git a/tests/training/test_dp_linear.py b/tests/training/test_dp_linear.py
new file mode 100644
index 000000000..0e0eae641
--- /dev/null
+++ b/tests/training/test_dp_linear.py
@@ -0,0 +1,47 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import pytest
+import torch
+from torch import nn
+
+from nemo_safe_synthesizer.privacy.dp_transformers import linear as linear_mod
+
+
+def test_compute_linear_grad_sample_returns_weight_and_bias_gradients():
+    layer = nn.Linear(3, 2)
+    activation = torch.tensor(
+        [
+            [1.0, 2.0, 3.0],
+            [4.0, 5.0, 6.0],
+        ]
+    )
+    backprops = torch.tensor(
+        [
+            [0.5, 1.5],
+            [2.0, 3.0],
+        ]
+    )
+
+    result = linear_mod.compute_linear_grad_sample(layer, [activation], backprops)
+
+    assert set(result) == {layer.weight, layer.bias}
+    torch.testing.assert_close(result[layer.weight], torch.einsum("ni,nj->nij", backprops, activation))
+    torch.testing.assert_close(result[layer.bias], backprops)
+
+
+def test_contract_tensor_rejects_non_tensor_result(monkeypatch: pytest.MonkeyPatch):
+    def fake_contract(*_args, **_kwargs) -> int:
+        return 1
+
+    monkeypatch.setattr(linear_mod, "contract", fake_contract)
+
+    with pytest.raises(TypeError, match=r"expected opt_einsum\.contract"):
+        linear_mod._contract_tensor("n->n", torch.ones(1))
+
+
+def test_linear_parameter_rejects_plain_tensor():
+    with pytest.raises(TypeError, match=r"expected nn\.Linear parameter"):
+        linear_mod._linear_parameter(torch.ones(1))