diff --git a/dpsynth/adapters/beam.py b/dpsynth/adapters/beam.py
index 933f69d4..8ce1bec7 100644
--- a/dpsynth/adapters/beam.py
+++ b/dpsynth/adapters/beam.py
@@ -24,7 +24,7 @@
 
 from __future__ import annotations
 
-from collections.abc import Callable
+from collections.abc import Callable, Mapping, Sequence
 import dataclasses
 import io
 import math
@@ -227,9 +227,9 @@ class _EncodeAndProject(beam.DoFn):
 
   def __init__(
       self,
-      column_measurements: dict[str, initialization.ColumnMeasurement],
-      domains: dict[str, Any],
-      workload: list[mbi.Clique],
+      column_measurements: Mapping[str, initialization.ColumnMeasurement],
+      domains: Mapping[str, Any],
+      workload: Sequence[mbi.Clique],
   ):
     super().__init__()
     # Reuse the shared per-column codec so Beam encoding matches the in-memory
@@ -298,9 +298,9 @@ class ComputeMarginals(beam.PTransform):
 
   def __init__(
       self,
-      column_measurements: dict[str, initialization.ColumnMeasurement],
-      domains: dict[str, Any],
-      workload: list[mbi.Clique],
+      column_measurements: Mapping[str, initialization.ColumnMeasurement],
+      domains: Mapping[str, Any],
+      workload: Sequence[mbi.Clique],
   ):
     super().__init__()
     self._column_measurements = column_measurements
diff --git a/dpsynth/adapters/pydantic_api.py b/dpsynth/adapters/pydantic_api.py
index f823fe73..a1fdf93c 100644
--- a/dpsynth/adapters/pydantic_api.py
+++ b/dpsynth/adapters/pydantic_api.py
@@ -14,6 +14,7 @@
 
 """Pydantic <--> DataFrame conversion utilities for TabularSynthesizer."""
 
+from collections.abc import Mapping
 import enum
 import inspect
 import math
@@ -141,7 +142,7 @@ def infer_domain_from_model(
 
 def models_to_dataframe(
     records: list[RecordT],
-    domains_dict: dict[str, domain.AttributeType],
+    domains_dict: Mapping[str, domain.AttributeType],
 ) -> pd.DataFrame:
   """Converts a list of pydantic models to a TabularSynthesizer-compatible DataFrame.
 
@@ -166,7 +167,7 @@ def models_to_dataframe(
 def dataframe_to_models(
     df: pd.DataFrame | data_generation_v3.DataGenerationResult,
     model_cls: type[RecordT],
-    domains_dict: dict[str, domain.AttributeType],
+    domains_dict: Mapping[str, domain.AttributeType],
 ) -> list[RecordT]:
   """Converts a synthetic DataFrame back to pydantic model instances.
 
diff --git a/dpsynth/discrete_mechanisms/common.py b/dpsynth/discrete_mechanisms/common.py
index 49eddcb1..8939df91 100644
--- a/dpsynth/discrete_mechanisms/common.py
+++ b/dpsynth/discrete_mechanisms/common.py
@@ -197,8 +197,8 @@ def exponential_mechanism(
 
 def measure_marginals_with_noise(
     rng: np.random.Generator,
-    data: mbi.Projectable,
-    marginal_queries: list[tuple[str, ...]],
+    data: mbi.Projectable | mbi.Dataset,
+    marginal_queries: Sequence[tuple[str, ...]],
     gdp_sigma: float,
     weights: np.ndarray | None = None,
     max_records_per_user: int = 1,
@@ -443,7 +443,7 @@ def score(cl):
 
 
 def compute_independence_errors(
-    data: mbi.Projectable,
+    data: mbi.Projectable | mbi.Dataset,
     model: mbi.MarkovRandomField,
     cliques: Sequence[mbi.Clique],
 ) -> dict[mbi.Clique, float]:
diff --git a/dpsynth/discrete_mechanisms/direct.py b/dpsynth/discrete_mechanisms/direct.py
index 37f6292c..a6b7ce13 100644
--- a/dpsynth/discrete_mechanisms/direct.py
+++ b/dpsynth/discrete_mechanisms/direct.py
@@ -39,7 +39,7 @@ def configure(self, _=None, *, zcdp_rho, delta=0, max_records_per_user=1):
 
   marginal_oracle: mbi.MarginalOracle | None = None
   pgm_iters: int = 5000
-  prespecified_marginal_queries: list[tuple[str, ...]] = dataclasses.field(
+  prespecified_marginal_queries: Sequence[tuple[str, ...]] = dataclasses.field(
       default_factory=list
   )
 
diff --git a/dpsynth/domain.py b/dpsynth/domain.py
index 6abe66f2..38355f09 100644
--- a/dpsynth/domain.py
+++ b/dpsynth/domain.py
@@ -58,7 +58,7 @@
 
 PathType = pathlib.Path
 
-CategoricalValue: TypeAlias = bool | int | str
+CategoricalValue: TypeAlias = bool | int | float | str
 
 IntervalHandling = Literal['midpoint', 'sample', 'interval']
 
@@ -193,7 +193,7 @@ class NumericalAttribute:
   min_value: float
   max_value: float
   clip_to_range: bool = True
-  sentinel: float | int | str | None = None
+  sentinel: float | int | str | np.integer | np.floating | None = None
   dtype: str = 'float'
   interval_handling: str = 'midpoint'
   description: str | None = None
@@ -232,7 +232,9 @@ def __post_init__(self):
         )
 
   @property
-  def resolved_sentinel(self) -> float | int | str:
+  def resolved_sentinel(
+      self,
+  ) -> float | int | str | np.integer | np.floating:
     """Returns the effective sentinel, with mode-appropriate defaults."""
     if self.sentinel is not None:
       return self.sentinel
diff --git a/dpsynth/local_mode/primitives.py b/dpsynth/local_mode/primitives.py
index 5d0e4c36..91a4cfe7 100644
--- a/dpsynth/local_mode/primitives.py
+++ b/dpsynth/local_mode/primitives.py
@@ -45,6 +45,7 @@
 
 from __future__ import annotations
 
+from typing import overload
 
 import numpy as np
 import scipy.stats
@@ -321,6 +322,36 @@ def _select_partitions_sips(
 # ---------------------------------------------------------------------------
 
 
+@overload
+def add_gaussian_noise(
+    rng: np.random.Generator,
+    counts: float | int,
+    sigma: float,
+    max_records_per_user: int = 1,
+) -> float:
+  ...
+
+
+@overload
+def add_gaussian_noise(
+    rng: np.random.Generator,
+    counts: np.ndarray,
+    sigma: float,
+    max_records_per_user: int = 1,
+) -> np.ndarray:
+  ...
+
+
+@overload
+def add_gaussian_noise(
+    rng: np.random.Generator,
+    counts: np.ndarray | float | int,
+    sigma: float,
+    max_records_per_user: int = 1,
+) -> float | np.ndarray:
+  ...
+
+
 def add_gaussian_noise(
     rng: np.random.Generator,
     counts: np.ndarray | float | int,
diff --git a/dpsynth/relational/domain.py b/dpsynth/relational/domain.py
index 3b42b69b..cf06b79e 100644
--- a/dpsynth/relational/domain.py
+++ b/dpsynth/relational/domain.py
@@ -232,7 +232,9 @@ def from_yaml_file(
 
 
 def to_dict(
-    table_domains: Mapping[str, domain.Schema],
+    table_domains: Mapping[
+        str, domain.Schema | Mapping[str, domain.AttributeType]
+    ],
     foreign_keys: Sequence[ForeignKeyRelation] = (),
 ) -> dict[str, Any]:
   """Converts multi-table schemas and foreign keys to a dictionary.
@@ -261,7 +263,9 @@ def to_dict(
 
 
 def to_yaml_file(
-    table_domains: Mapping[str, domain.Schema],
+    table_domains: Mapping[
+        str, domain.Schema | Mapping[str, domain.AttributeType]
+    ],
     foreign_keys: Sequence[ForeignKeyRelation],
     filepath: str | PathType,
 ) -> None:
diff --git a/dpsynth/relational/synthesizer.py b/dpsynth/relational/synthesizer.py
index 29c50c4f..9de342d7 100644
--- a/dpsynth/relational/synthesizer.py
+++ b/dpsynth/relational/synthesizer.py
@@ -19,7 +19,7 @@
 from collections.abc import Collection, Hashable, Mapping, Sequence
 import dataclasses
 import math
-from typing import Any, Literal
+from typing import Any, Literal, TypeAlias
 
 from absl import logging
 import dp_accounting
@@ -42,9 +42,13 @@
 _LOGGING_UNUSED = logging
 # pylint: enable=unused-import
 
+TableDomains: TypeAlias = Mapping[
+    str, domain.Schema | Mapping[str, domain.AttributeType]
+]
+
 
 def _validate_input_table_columns(
-    domains: Mapping[str, domain.Schema],
+    domains: TableDomains,
     foreign_keys: Sequence[rel_domain.ForeignKeyRelation],
     table_columns: Mapping[str, Collection[str]],
 ) -> None:
@@ -264,7 +268,7 @@ def _run_table_initializers(
 
 
 def _encode_and_compress_tables(
-    domains: Mapping[str, domain.Schema],
+    domains: TableDomains,
     table_measurements: Mapping[
         str, Mapping[str, initialization.ColumnMeasurement]
     ],
@@ -419,7 +423,7 @@ def _run_table_preprocessing(
 
 
 def _create_table_initializers(
-    domains: Mapping[str, domain.Schema],
+    domains: TableDomains,
     numerical_bins: int,
 ) -> dict[str, dict[str, api.MechanismConfig]]:
   """Creates per-table and per-column initializers from relational schemas."""
@@ -430,7 +434,7 @@ def _create_table_initializers(
 
 
 def _compute_table_col_deltas(
-    domains: Mapping[str, domain.Schema],
+    domains: TableDomains,
     delta: float,
     init_budget_fraction: float,
 ) -> dict[str, dict[str, float]]:
@@ -847,7 +851,7 @@ def _decompress_synthetic_datasets(
 def _decode_synthetic_tables(
     decompressed_datasets: Mapping[str, mbi.Dataset],
     column_codecs: Mapping[str, data_generation_v3.TabularCodec],
-    domains: Mapping[str, domain.Schema],
+    domains: TableDomains,
     rng: np.random.Generator,
 ) -> dict[str, pd.DataFrame]:
   """Decodes discrete datasets into continuous/categorical DataFrames.
diff --git a/dpsynth/text/bulk_inference.py b/dpsynth/text/bulk_inference.py
index 4cf28ea2..8fa5715a 100644
--- a/dpsynth/text/bulk_inference.py
+++ b/dpsynth/text/bulk_inference.py
@@ -23,7 +23,7 @@
 import random
 import re
 import time
-from typing import Protocol, TypeVar
+from typing import Any, Protocol, TypeVar
 
 from absl import logging
 from dpsynth import domain
@@ -376,7 +376,7 @@ def _categorical_json_type(
 
 def domain_to_json_schema(
     domain_spec: Mapping[str, domain.AttributeType],
-) -> dict[str, object]:
+) -> dict[str, Any]:
   """Converts a dpsynth Domain to a JSON schema dict for GenAI."""
   properties = {}
   for name, attr in domain_spec.items():
diff --git a/dpsynth/transformations.py b/dpsynth/transformations.py
index c4ca6a91..0be086f1 100644
--- a/dpsynth/transformations.py
+++ b/dpsynth/transformations.py
@@ -16,20 +16,22 @@
 
 import bisect
 from collections.abc import Callable, Mapping, Sequence
+import dataclasses
 import math
 from typing import Any, Generic, TypeAlias, TypeVar
 
-import attr
 from dpsynth import domain
 import numpy as np
 import pandas as pd
 
 CategoricalValue: TypeAlias = bool | int | float | str
-R, T, S = TypeVar('R'), TypeVar('T'), TypeVar('S')
+R = TypeVar('R')
+T = TypeVar('T')
+S = TypeVar('S')
 
 
-@attr.define(frozen=True)
-class DataTransformation(Generic[R, T]):  # pyrefly: ignore[not-a-type]
+@dataclasses.dataclass(frozen=True)
+class DataTransformation(Generic[R, T]):
   """Dataclass for transforming data from one domain to another.
 
   DataTransformations are both reversible (via inverse) and composable (via @).
@@ -48,22 +50,22 @@ class DataTransformation(Generic[R, T]):  # pyrefly: ignore[not-a-type]
   1
   """
 
-  transform: Callable[[R], T] | Mapping[R, T] = attr.field()  # pyrefly: ignore[not-a-type]
-  inverse_transform: Callable[[T], R] | Mapping[T, R] = attr.field()  # pyrefly: ignore[not-a-type]
+  transform: Callable[[R], T] | Mapping[R, T]
+  inverse_transform: Callable[[T], R] | Mapping[T, R]
 
-  def __call__(self, value: R) -> T:  # pyrefly: ignore[not-a-type]
+  def __call__(self, value: R) -> T:
     if isinstance(self.transform, Mapping):
       return self.transform[value]
     return self.transform(value)
 
   @property
-  def inverse(self) -> 'DataTransformation[T, R]':  # pyrefly: ignore[not-a-type]
+  def inverse(self) -> 'DataTransformation[T, R]':
     """The reverse transformation of this instance."""
-    return DataTransformation(self.inverse_transform, self.transform)  # pyrefly: ignore[bad-argument-count]
+    return DataTransformation(self.inverse_transform, self.transform)
 
   def __matmul__(
-      self, other: 'DataTransformation[T, S]'  # pyrefly: ignore[not-a-type]
-  ) -> 'DataTransformation[R, S]':  # pyrefly: ignore[not-a-type]
+      self, other: 'DataTransformation[S, R]'
+  ) -> 'DataTransformation[S, T]':
     """Returns a DataTransformation that composes this instance with other.
 
     Example Usage:
@@ -82,8 +84,8 @@ def __matmul__(
       A DataTransformation that composes this instance with other.
     """
     return DataTransformation(
-        lambda x: self(other(x)),  # pyrefly: ignore[bad-argument-count]
-        lambda x: other.inverse(self.inverse(x)),
+        lambda x: self(other(x)),  # pyrefly: ignore[bad-argument-type]
+        lambda x: other.inverse(self.inverse(x)),  # pyrefly: ignore[bad-argument-type]
     )
 
 
@@ -129,10 +131,10 @@ def discrete_encoder(
   ood = attribute_domain.out_of_domain_index
   transform = lambda v: index_map.get(value_type(v), ood)
   reverse = dict(enumerate(attribute_domain.possible_values))
-  return DataTransformation(transform, reverse)  # pyrefly: ignore[bad-argument-count]
+  return DataTransformation(transform, reverse)
 
 
-@attr.define(frozen=True)
+@dataclasses.dataclass(frozen=True)
 class _Interval:
   """A numeric interval with a string representation."""
 
@@ -199,7 +201,7 @@ def create_discretize_transformation(
       attribute_domain.max_value,
   ]
   intervals = [
-      _Interval(left, right, closed_left=(i == 0))  # pyrefly: ignore[bad-argument-count, unexpected-keyword]
+      _Interval(left, right, closed_left=(i == 0))
       for i, (left, right) in enumerate(zip(bin_edges[:-1], bin_edges[1:]))
   ]
   interval_strs = [str(iv) for iv in intervals]
@@ -216,7 +218,7 @@ def transform(value: Any) -> str:
     idx = bisect.bisect_left(inner_edges, value)
     return interval_strs[idx]
 
-  def reverse(value: str) -> float | str:
+  def reverse(value: str) -> float | int | str | np.integer | np.floating:
     if value == ood_sentinel:
       return sentinel
     idx = interval_strs.index(value)
@@ -231,8 +233,8 @@ def reverse(value: str) -> float | str:
       return math.ceil(result)
     return result
 
-  new_domain = domain.CategoricalAttribute(possible_values)  # pyrefly: ignore[bad-argument-count]
-  transformation = DiscretizeTransformation(transform, reverse)  # pyrefly: ignore[bad-argument-count]
+  new_domain = domain.CategoricalAttribute(possible_values)
+  transformation = DiscretizeTransformation(transform, reverse)
   return new_domain, transformation
 
 
diff --git a/tests/adapters/beam_test.py b/tests/adapters/beam_test.py
index 78e5172f..99bfe462 100644
--- a/tests/adapters/beam_test.py
+++ b/tests/adapters/beam_test.py
@@ -18,6 +18,7 @@
 import multiprocessing
 import os
 import tempfile
+from typing import Any, cast
 from unittest import mock
 
 from absl.testing import absltest
@@ -34,7 +35,7 @@
 import numpy as np
 
 _manager = None
-_test_results = None
+_test_results: Any = None
 
 
 def setUpModule():
@@ -491,8 +492,8 @@ def test_configure_returns_calibrated_wrapper(self):
     # dp_event is delegated to the wrapped, now-calibrated synthesizer.
     self.assertIsNotNone(configured.dp_event)
     # The original wrapper is left uncalibrated (configure returns a copy).
-    with self.assertRaises(Exception):
-      _ = beam_synth.dp_event
+    with self.assertRaises(AttributeError):
+      _ = getattr(beam_synth, 'dp_event')
 
   def test_inherited_calibrate_produces_calibrated_wrapper(self):
     # calibrate is inherited from DPMechanism; it binary-searches a zCDP budget
@@ -510,8 +511,8 @@ def test_uncalibrated_call_raises(self):
     beam_synth = beam_adapter.BeamTabularConfig(
         data_generation_v3.TabularConfig()
     )
-    with self.assertRaises(Exception):
-      beam_synth(np.random.default_rng(0), lambda p: p)
+    with self.assertRaises(TypeError):
+      cast(Any, beam_synth)(np.random.default_rng(0), lambda p: p)
 
   def test_honors_temp_location(self):
     domains = {'a': domain.CategoricalAttribute(possible_values=['x', 'y'])}
diff --git a/tests/data_generation_test.py b/tests/data_generation_test.py
index 2456cad0..5a701062 100644
--- a/tests/data_generation_test.py
+++ b/tests/data_generation_test.py
@@ -187,6 +187,7 @@ def test_format_output_data(self, output_format, expected_output):
     )
     input_data = [(1,), (2,)]
 
+    assert config.output_format is not None
     backend = pipeline_dp.LocalBackend()
     formatted_data = list(
         data_generation._format_output_data(
@@ -357,6 +358,7 @@ def test_generate_from_model_e2e(self):
         attribute="uncompress",
         side_effect=lambda x: x,
     ):
+      assert config.output_format is not None
       output_data = data_generation.generate_from_model(
           models, descriptors, 10, config.output_format
       )
@@ -410,6 +412,7 @@ def test_generate_from_model_tfrecord_e2e(self):
         attribute="uncompress",
         side_effect=lambda x: x,
     ):
+      assert config.output_format is not None
       output_data = data_generation.generate_from_model(
           models, descriptors, 10, config.output_format
       )
diff --git a/tests/data_generation_v3_test.py b/tests/data_generation_v3_test.py
index f5e8c961..759a1dae 100644
--- a/tests/data_generation_v3_test.py
+++ b/tests/data_generation_v3_test.py
@@ -14,6 +14,8 @@
 
 from __future__ import annotations
 
+from typing import Any, cast
+
 from absl.testing import absltest
 from absl.testing import parameterized
 import dp_accounting
@@ -31,7 +33,7 @@
 
 
 def _make_discrete_data(rng, n=1000):
-  domains = mbi.Domain(['a', 'b', 'c'], [3, 3, 3])
+  domains = mbi.Domain(('a', 'b', 'c'), (3, 3, 3))
   a = rng.integers(0, 3, size=n)
   b = np.where(rng.random(n) < 0.75, a, rng.integers(0, 3, size=n))
   c = (a + b + rng.integers(0, 2, size=n)) % 3
@@ -196,7 +198,7 @@ def test_raises_when_not_calibrated(self):
     rng = np.random.default_rng(0)
     v3 = TabularConfig()
     with self.assertRaises(Exception):
-      v3(rng, df)
+      cast(Any, v3)(rng, df)
 
   def test_dp_event_returns_composed_event(self):
     domains = {
@@ -378,7 +380,9 @@ def test_configure_propagates_k_to_submechanisms(self):
         self._categorical_domains(), zcdp_rho=100.0, max_records_per_user=k
     )
     self.assertEqual(calibrated.max_records_per_user, k)
-    self.assertEqual(calibrated.base_mechanism.max_records_per_user, k)
+    self.assertEqual(
+        getattr(calibrated.base_mechanism, 'max_records_per_user'), k
+    )
 
   def test_dp_event_invariant_to_k(self):
     config = TabularConfig()
diff --git a/tests/dataset_descriptors/dataset_descriptor_test.py b/tests/dataset_descriptors/dataset_descriptor_test.py
index db43edf4..0fec7aa6 100644
--- a/tests/dataset_descriptors/dataset_descriptor_test.py
+++ b/tests/dataset_descriptors/dataset_descriptor_test.py
@@ -119,8 +119,8 @@ def test_getstate(self):
             possible_values=[1, 2, 3]
         ),
     )
-    attr_desc._encoding_transform = mock.Mock()
-    attr_desc._compress_transform = mock.Mock()
+    attr_desc.__dict__["encoding_transform"] = mock.Mock()
+    attr_desc.__dict__["compress_transform"] = mock.Mock()
 
     state = attr_desc.__getstate__()
     self.assertNotIn("encoding_transform", state)
diff --git a/tests/discrete_mechanisms/aim_test.py b/tests/discrete_mechanisms/aim_test.py
index 7370d853..1478e4b3 100644
--- a/tests/discrete_mechanisms/aim_test.py
+++ b/tests/discrete_mechanisms/aim_test.py
@@ -22,7 +22,7 @@
 
 
 def _make_correlated_dataset(rng, n=1000):
-  domain = mbi.Domain(["a", "b", "c"], [3, 3, 3])
+  domain = mbi.Domain(("a", "b", "c"), (3, 3, 3))
   a = rng.integers(0, 3, size=n)
   b = np.where(rng.random(n) < 0.75, a, rng.integers(0, 3, size=n))
   c = (a + b + rng.integers(0, 2, size=n)) % 3
@@ -59,7 +59,7 @@ def _correlated_workload_mechanism_baseline_errors(
 class AIMTest(absltest.TestCase):
 
   def test_fits_one_way_marginals_with_aim(self):
-    data = mbi.Dataset.synthetic(mbi.Domain(["a", "b", "c"], [3, 4, 5]), N=1000)
+    data = mbi.Dataset.synthetic(mbi.Domain(("a", "b", "c"), (3, 4, 5)), N=1000)
     workload = [("a",), ("b",), ("c",)]
     config = aim.AIMConfig(workload=workload, max_rounds=4, pgm_iters=500)
 
@@ -74,7 +74,7 @@ def test_fits_one_way_marginals_with_aim(self):
       np.testing.assert_allclose(actual, expected, atol=1)
 
   def test_fits_one_way_marginals_with_aim_gdp(self):
-    data = mbi.Dataset.synthetic(mbi.Domain(["a", "b", "c"], [3, 4, 5]), N=1000)
+    data = mbi.Dataset.synthetic(mbi.Domain(("a", "b", "c"), (3, 4, 5)), N=1000)
     workload = [("a",), ("b",), ("c",)]
 
     config = aim_gdp.AIMGDPConfig(
diff --git a/tests/discrete_mechanisms/common_test.py b/tests/discrete_mechanisms/common_test.py
index 23c61eb0..7b66f8eb 100644
--- a/tests/discrete_mechanisms/common_test.py
+++ b/tests/discrete_mechanisms/common_test.py
@@ -45,7 +45,7 @@ def test_exponential_mechanism(self):
     self.assertEqual(idx, 1)
 
   def test_measure_marginals_with_noise(self):
-    data = mbi.Dataset.synthetic(mbi.Domain(["a", "b", "c"], [3, 4, 5]), N=1000)
+    data = mbi.Dataset.synthetic(mbi.Domain(("a", "b", "c"), (3, 4, 5)), N=1000)
     marginal_queries = [("a",), ("b",), ("c",)]
     measurements = common.measure_marginals_with_noise(
         np.random.default_rng(0), data, marginal_queries, gdp_sigma=1.0
@@ -75,7 +75,7 @@ def test_get_domain_compression_transformations(self):
     self.assertEqual(compressed_domain, mbi.Domain.fromdict({"a": 5}))
 
   def test_supporting_cliques(self):
-    domain = mbi.Domain(["a", "b", "c", "d"], [3, 3, 3, 100])
+    domain = mbi.Domain(("a", "b", "c", "d"), (3, 3, 3, 100))
     cliques = common.supporting_cliques(domain, workload=None)
     self.assertCountEqual(
         cliques, list(itertools.combinations(domain.attributes, 3))
diff --git a/tests/discrete_mechanisms/direct_test.py b/tests/discrete_mechanisms/direct_test.py
index 0ecf3ed4..d39466b3 100644
--- a/tests/discrete_mechanisms/direct_test.py
+++ b/tests/discrete_mechanisms/direct_test.py
@@ -22,7 +22,7 @@
 class DirectTest(absltest.TestCase):
 
   def test_fits_one_way_marginals(self):
-    data = mbi.Dataset.synthetic(mbi.Domain(['a', 'b', 'c'], [3, 4, 5]), N=1000)
+    data = mbi.Dataset.synthetic(mbi.Domain(('a', 'b', 'c'), (3, 4, 5)), N=1000)
 
     prespecified_queries = [('a', 'b'), ('a', 'c'), ('b', 'c')]
     config = direct.DirectConfig(
diff --git a/tests/discrete_mechanisms/discrete_mechanisms_test.py b/tests/discrete_mechanisms/discrete_mechanisms_test.py
index 2e1e86a2..819479d0 100644
--- a/tests/discrete_mechanisms/discrete_mechanisms_test.py
+++ b/tests/discrete_mechanisms/discrete_mechanisms_test.py
@@ -52,7 +52,7 @@
 
 def _make_skewed_dataset(rng):
   """Creates a dataset where column 'a' concentrates in 3 of 10 bins."""
-  domain = mbi.Domain(['a', 'b', 'c'], [10, 4, 5])
+  domain = mbi.Domain(('a', 'b', 'c'), (10, 4, 5))
   df = {col: rng.integers(0, domain[col], size=1000) for col in domain}
   df['a'] = rng.choice(3, size=1000)  # Only bins 0-2 populated.
   return mbi.Dataset(df, domain)
@@ -63,14 +63,16 @@ class SupportingCliquesSufficiencyTest(parameterized.TestCase):
 
   @parameterized.named_parameters(*_MECHANISMS.items())
   def test_mechanism_runs_on_precomputed_marginals(self, mechanism):
-    domain = mbi.Domain(['a', 'b', 'c', 'd'], [3, 4, 5, 6])
+    domain = mbi.Domain(('a', 'b', 'c', 'd'), (3, 4, 5, 6))
     data = mbi.Dataset.synthetic(domain, N=500)
     rng = np.random.default_rng(42)
 
     calibrated = mechanism.configure(zcdp_rho=_ZCDP_RHO)
     cliques = mechanism.supporting_cliques(domain)
 
-    precomputed = mbi.CliqueVector.from_projectable(data, cliques)
+    precomputed = mbi.CliqueVector.from_projectable(
+        data, cliques  # pyrefly: ignore[bad-argument-type]
+    )
 
     result = calibrated(rng, precomputed)
     self.assertIsInstance(result, common.DiscreteMechanismResult)
diff --git a/tests/discrete_mechanisms/discrete_test.py b/tests/discrete_mechanisms/discrete_test.py
index 3098d92b..10e1ea96 100644
--- a/tests/discrete_mechanisms/discrete_test.py
+++ b/tests/discrete_mechanisms/discrete_test.py
@@ -52,7 +52,7 @@ def test_configure_zero_one_way_fraction(self):
   def test_supporting_cliques_delegates(self):
     inner = MSTConfig(pgm_iters=500)
     config = DiscreteConfig(mechanism=inner)
-    domain = mbi.Domain(['a', 'b', 'c'], [3, 4, 5])
+    domain = mbi.Domain(('a', 'b', 'c'), (3, 4, 5))
     self.assertEqual(
         config.supporting_cliques(domain),
         inner.supporting_cliques(domain),
@@ -62,7 +62,7 @@ def test_supporting_cliques_delegates(self):
 class DiscreteMechanismTest(absltest.TestCase):
 
   def test_full_pipeline(self):
-    domain = mbi.Domain(['a', 'b', 'c'], [3, 4, 5])
+    domain = mbi.Domain(('a', 'b', 'c'), (3, 4, 5))
     data = mbi.Dataset.synthetic(domain, N=500)
     rng = np.random.default_rng(42)
 
@@ -76,7 +76,7 @@ def test_full_pipeline(self):
     self.assertEqual(result.synthetic_data.domain, domain)
 
   def test_with_initial_measurements_skips_one_way(self):
-    domain = mbi.Domain(['a', 'b', 'c'], [3, 4, 5])
+    domain = mbi.Domain(('a', 'b', 'c'), (3, 4, 5))
     data = mbi.Dataset.synthetic(domain, N=500)
     rng = np.random.default_rng(42)
 
@@ -92,7 +92,7 @@ def test_with_initial_measurements_skips_one_way(self):
     self.assertIsInstance(result, common.DiscreteMechanismResult)
 
   def test_compression_restores_domain(self):
-    domain = mbi.Domain(['a', 'b', 'c'], [10, 4, 5])
+    domain = mbi.Domain(('a', 'b', 'c'), (10, 4, 5))
     rng = np.random.default_rng(0)
     df = {col: rng.integers(0, domain[col], size=1000) for col in domain}
     df['a'] = rng.choice(3, size=1000)
@@ -120,7 +120,7 @@ def test_calibrate_works(self):
     config = DiscreteConfig(
         mechanism=MSTConfig(pgm_iters=500),
     )
-    domain = mbi.Domain(['a', 'b'], [3, 4])
+    domain = mbi.Domain(('a', 'b'), (3, 4))
     data = mbi.Dataset.synthetic(domain, N=200)
     rng = np.random.default_rng(0)
 
diff --git a/tests/discrete_mechanisms/independent_test.py b/tests/discrete_mechanisms/independent_test.py
index 74dba1be..fc96e142 100644
--- a/tests/discrete_mechanisms/independent_test.py
+++ b/tests/discrete_mechanisms/independent_test.py
@@ -23,7 +23,7 @@ class IndependentTest(absltest.TestCase):
 
   def test_fits_one_way_marginals(self):
     """Independent with externally-supplied 1-ways should recover marginals."""
-    data = mbi.Dataset.synthetic(mbi.Domain(['a', 'b', 'c'], [3, 4, 5]), N=1000)
+    data = mbi.Dataset.synthetic(mbi.Domain(('a', 'b', 'c'), (3, 4, 5)), N=1000)
 
     config = independent.IndependentConfig()
 
@@ -56,7 +56,7 @@ def test_fits_one_way_marginals(self):
 
   def test_skips_duplicate_cliques_from_initial_measurements(self):
     """Independent should not re-measure pre-measured cliques."""
-    data = mbi.Dataset.synthetic(mbi.Domain(['a', 'b', 'c'], [3, 4, 5]), N=100)
+    data = mbi.Dataset.synthetic(mbi.Domain(('a', 'b', 'c'), (3, 4, 5)), N=100)
     # Pre-measure column 'a'.
     marginal_a = data.project(('a',)).datavector()
     initial = [mbi.LinearMeasurement(marginal_a, ('a',), stddev=1.0)]
diff --git a/tests/discrete_mechanisms/mst_test.py b/tests/discrete_mechanisms/mst_test.py
index 4ddb1727..5a650ddf 100644
--- a/tests/discrete_mechanisms/mst_test.py
+++ b/tests/discrete_mechanisms/mst_test.py
@@ -76,7 +76,7 @@ def test_dp_maximum_spanning_tree_infinite_eps(self):
 
   def test_fits_one_way_marginals(self):
     """MST + externally-supplied 1-ways should recover all one-way marginals."""
-    data = mbi.Dataset.synthetic(mbi.Domain(['a', 'b', 'c'], [3, 4, 5]), N=1000)
+    data = mbi.Dataset.synthetic(mbi.Domain(('a', 'b', 'c'), (3, 4, 5)), N=1000)
 
     calibrated = mst.MSTConfig(pgm_iters=500).configure(zcdp_rho=10000)
 
diff --git a/tests/discrete_mechanisms/swift_test.py b/tests/discrete_mechanisms/swift_test.py
index b11afb47..1d3ec007 100644
--- a/tests/discrete_mechanisms/swift_test.py
+++ b/tests/discrete_mechanisms/swift_test.py
@@ -28,7 +28,7 @@ class CliqueTreeTest(absltest.TestCase):
 
   def setUp(self):
     super().setUp()
-    self.domain = mbi.Domain(['a', 'b', 'c', 'd'], [2, 3, 4, 5])
+    self.domain = mbi.Domain(('a', 'b', 'c', 'd'), (2, 3, 4, 5))
 
   def test_best_supporting_edge(self):
     edges = [(('a',), ('b',)), (('b',), ('c',))]
@@ -77,9 +77,13 @@ def test_local_update(self):
 class SWIFTTest(absltest.TestCase):
 
   def test_select_queries(self):
-    errors = {('a', 'b'): 100, ('b', 'c'): 100, ('a', 'c'): 100}
-    domain = mbi.Domain(['a', 'b', 'c'], [2, 3, 4])
-    candidates = {key: 1.0 for key in errors}
+    errors: dict[mbi.Clique, float] = {
+        ('a', 'b'): 100.0,
+        ('b', 'c'): 100.0,
+        ('a', 'c'): 100.0,
+    }
+    domain = mbi.Domain(('a', 'b', 'c'), (2, 3, 4))
+    candidates: dict[mbi.Clique, float] = {key: 1.0 for key in errors}
     max_clique_size = 100
     gdp_budget = 100.0
     selected, jtree = swift.select_queries(
@@ -97,9 +101,13 @@ def test_select_queries(self):
     self.assertAlmostEqual(selected[('b', 'c')], expected_budget_bc)
 
   def test_select_queries_nonpositive_errors(self):
-    errors = {('a', 'b'): 0.0, ('b', 'c'): -1.0, ('a', 'c'): -2.0}
-    domain = mbi.Domain(['a', 'b', 'c'], [2, 3, 4])
-    candidates = {key: 1.0 for key in errors}
+    errors: dict[mbi.Clique, float] = {
+        ('a', 'b'): 0.0,
+        ('b', 'c'): -1.0,
+        ('a', 'c'): -2.0,
+    }
+    domain = mbi.Domain(('a', 'b', 'c'), (2, 3, 4))
+    candidates: dict[mbi.Clique, float] = {key: 1.0 for key in errors}
     gdp_budget = 100.0
 
     selected, jtree = swift.select_queries(
@@ -120,9 +128,11 @@ def test_best_subset_and_allocation(self):
     self.assertLen(allocation, 3)
 
   def test_build_clique_tree(self):
-    domain = mbi.Domain(['a', 'b', 'c', 'd', 'e', 'f'], [3, 4, 5, 6, 7, 8])
+    domain = mbi.Domain(('a', 'b', 'c', 'd', 'e', 'f'), (3, 4, 5, 6, 7, 8))
     max_clique_size = 100
-    errors = {key: 1.0 for key in itertools.combinations(domain.attributes, 2)}
+    errors: dict[mbi.Clique, float] = {
+        key: 1.0 for key in itertools.combinations(domain.attributes, 2)
+    }
 
     tree = swift.build_clique_tree(domain, errors, max_clique_size)
     actual_max_clique_size = max(domain.size(cl) for cl in tree.nodes)
@@ -133,7 +143,7 @@ def test_build_clique_tree(self):
     self.assertLessEqual(actual_max_clique_size, max_clique_size)
 
   def test_fits_one_way_marginals(self):
-    data = mbi.Dataset.synthetic(mbi.Domain(['a', 'b', 'c'], [3, 4, 5]), N=1000)
+    data = mbi.Dataset.synthetic(mbi.Domain(('a', 'b', 'c'), (3, 4, 5)), N=1000)
 
     config = swift.SWIFTConfig(pgm_iters=500).configure(zcdp_rho=10000)
 
diff --git a/tests/domain_test.py b/tests/domain_test.py
index a5e740bd..4e2a6b66 100644
--- a/tests/domain_test.py
+++ b/tests/domain_test.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import math
+from typing import Any, cast
 
 from absl.testing import absltest
 from dpsynth import domain
@@ -48,7 +49,7 @@ def test_mixed_types_rejected(self):
     with self.assertRaises(ValueError):
       domain.CategoricalAttribute(possible_values=['a', 1])
     with self.assertRaises(ValueError):
-      domain.CategoricalAttribute(possible_values=[None, 'a'])
+      domain.CategoricalAttribute(possible_values=cast(Any, [None, 'a']))
 
   def test_invalid_range(self):
     with self.assertRaises(ValueError):
@@ -126,7 +127,9 @@ def test_sentinel_yaml_roundtrip(self):
     temp_file = self.create_tempfile('temp.yaml', mode='w+')
     domain.to_yaml_file(original, temp_file.full_path)
     loaded = domain.from_yaml_file(temp_file.full_path)
-    self.assertEqual(loaded['num'].sentinel, -1)
+    loaded_num = loaded['num']
+    assert isinstance(loaded_num, domain.NumericalAttribute)
+    self.assertEqual(loaded_num.sentinel, -1)
 
   def test_string_sentinel_allowed_with_interval_handling(self):
     attr = domain.NumericalAttribute(
@@ -154,7 +157,8 @@ def test_numpy_numeric_sentinel_accepted(self):
     attr = domain.NumericalAttribute(0, 10, sentinel=np.int32(-1))
     self.assertEqual(attr.sentinel, -1)
     attr = domain.NumericalAttribute(0, 10, sentinel=np.float32(0.5))
-    self.assertAlmostEqual(attr.sentinel, 0.5, places=5)
+    assert isinstance(attr.sentinel, (float, np.floating))
+    self.assertAlmostEqual(float(attr.sentinel), 0.5, places=5)
 
   def test_freeform_text_defaults(self):
     attribute = domain.FreeFormTextAttribute()
diff --git a/tests/pipeline_transformations/categorical_values_derivation_test.py b/tests/pipeline_transformations/categorical_values_derivation_test.py
index a2ff7ef7..4cca9999 100644
--- a/tests/pipeline_transformations/categorical_values_derivation_test.py
+++ b/tests/pipeline_transformations/categorical_values_derivation_test.py
@@ -37,6 +37,8 @@ def test_derive_categorical_values(self):
         attribute_keys_to_derive=[0, 2],
     )
     accountant.compute_budgets()
+    self.assertIsNotNone(got)
+    assert got is not None
     got = list(got)
     self.assertEqual(
         got[0],
diff --git a/tests/pipeline_transformations/dataset_compression_test.py b/tests/pipeline_transformations/dataset_compression_test.py
index 85b9c1f0..43ac61f8 100644
--- a/tests/pipeline_transformations/dataset_compression_test.py
+++ b/tests/pipeline_transformations/dataset_compression_test.py
@@ -27,7 +27,9 @@ class FakeDataRecordConverter(dataset_descriptor.DataRecordConverter):
   def to_tuple(self, record: Any) -> tuple[Any, ...]:
     raise NotImplementedError("to_tuple is not implemented")
 
-  def from_tuple(self, record: tuple[Any, ...]) -> Any:
+  def from_tuple(
+      self, record: tuple[Any, ...], proto_object: Any | None = None
+  ) -> Any:
     raise NotImplementedError("from_tuple is not implemented")
 
 
diff --git a/tests/pipeline_transformations/dataset_encoding_test.py b/tests/pipeline_transformations/dataset_encoding_test.py
index 8b6d2c99..eeb65b77 100644
--- a/tests/pipeline_transformations/dataset_encoding_test.py
+++ b/tests/pipeline_transformations/dataset_encoding_test.py
@@ -30,7 +30,9 @@ class FakeDataRecordConverter(dataset_descriptor.DataRecordConverter):
   def to_tuple(self, record: Any) -> tuple[Any, ...]:
     raise NotImplementedError("to_tuple is not implemented")
 
-  def from_tuple(self, record: tuple[Any, ...]) -> Any:
+  def from_tuple(
+      self, record: tuple[Any, ...], proto_object: Any | None = None
+  ) -> Any:
     raise NotImplementedError("from_tuple is not implemented")
 
 
diff --git a/tests/pipeline_transformations/input_output_test.py b/tests/pipeline_transformations/input_output_test.py
index f491c04e..50763b63 100644
--- a/tests/pipeline_transformations/input_output_test.py
+++ b/tests/pipeline_transformations/input_output_test.py
@@ -144,7 +144,7 @@ def test_save_load_model_local(self):
     )
     domain_mrf = mbi.Domain(attributes=(0,), shape=(2,))
     clique = (0,)
-    factor = mbi.Factor(domain=domain_mrf, values=np.array([1, 2]))
+    factor = mbi.Factor(domain=domain_mrf, values=jnp.array([1, 2]))
     clique_vector = mbi.CliqueVector(
         domain_mrf,
         [clique],
@@ -177,7 +177,7 @@ def test_save_load_model_pipeline(self):
     )
     domain_mrf = mbi.Domain(attributes=(0,), shape=(2,))
     clique = (0,)
-    factor = mbi.Factor(domain=domain_mrf, values=np.array([1, 2]))
+    factor = mbi.Factor(domain=domain_mrf, values=jnp.array([1, 2]))
     clique_vector = mbi.CliqueVector(
         domain_mrf,
         [clique],
diff --git a/tests/pipeline_transformations/marginals_computations_test.py b/tests/pipeline_transformations/marginals_computations_test.py
index 08d3751b..1f1fc1e5 100644
--- a/tests/pipeline_transformations/marginals_computations_test.py
+++ b/tests/pipeline_transformations/marginals_computations_test.py
@@ -29,7 +29,7 @@ def test_beam_backend(self):
 
       queries = pipeline | "Create queries" >> beam.Create([[(0, 1), (2,)]])
       domain = pipeline | "Create domain" >> beam.Create(
-          [mbi.Domain([0, 1, 2], [2, 3, 4])]
+          [mbi.Domain((0, 1, 2), (2, 3, 4))]
       )
       data = pipeline | "Create" >> beam.Create(data)
       backend = pipeline_dp.BeamBackend()
@@ -45,7 +45,7 @@ def test_compute_exact_marginals(self):
     data = [(0, 1, 2), (0, 0, 2), (0, 2, 2), (1, 1, 2)]
 
     queries = [[(0, 1), (2,)]]  # singleton list
-    domain = [mbi.Domain([0, 1, 2], [2, 3, 4])]
+    domain = [mbi.Domain((0, 1, 2), (2, 3, 4))]
     marginals = dict(
         list(
             marginals_computations.compute_exact_marginals(
@@ -106,7 +106,7 @@ def create_dp_engine(self):
   def test_one_way_dp_marginals_empty_input(self):
     backend = pipeline_dp.LocalBackend()
     data = []
-    domain = [mbi.Domain(["a", "b", "c", "d"], [3, 4, 5, 6])]
+    domain = [mbi.Domain(("a", "b", "c", "d"), (3, 4, 5, 6))]
     dp_engine, accountant = self.create_dp_engine()
     result = marginals_computations.compute_one_way_dp_marginals(
         backend, dp_engine, data, domain, 1
@@ -121,7 +121,7 @@ def test_one_way_dp_marginals_single_row(self):
     backend = pipeline_dp.LocalBackend()
     dp_engine, accountant = self.create_dp_engine()
     data = [(0, 1)]
-    domain = [mbi.Domain(["col0", "col1"], [2, 2])]
+    domain = [mbi.Domain(("col0", "col1"), (2, 2))]
     result = marginals_computations.compute_one_way_dp_marginals(
         backend, dp_engine, data, domain, 2
     )
@@ -136,7 +136,7 @@ def test_one_way_dp_marginals_multiple_rows_different_values(self):
     backend = pipeline_dp.LocalBackend()
     dp_engine, accountant = self.create_dp_engine()
     data = [(0, 1, 2), (1, 0, 3)]
-    domain = [mbi.Domain(["col0", "col1", "col2"], [3, 4, 5])]
+    domain = [mbi.Domain(("col0", "col1", "col2"), (3, 4, 5))]
     result = marginals_computations.compute_one_way_dp_marginals(
         backend, dp_engine, data, domain, 3
     )
@@ -156,7 +156,7 @@ def test_one_way_dp_marginals_multiple_rows_repeated_values(self):
     backend = pipeline_dp.LocalBackend()
     dp_engine, accountant = self.create_dp_engine()
     data = [(0, 2), (0, 2), (2, 0), (0, 2)]
-    domain = [mbi.Domain(["col0", "col1"], [3, 3])]
+    domain = [mbi.Domain(("col0", "col1"), (3, 3))]
     result = marginals_computations.compute_one_way_dp_marginals(
         backend, dp_engine, data, domain, 2
     )
diff --git a/tests/pipeline_transformations/model_test.py b/tests/pipeline_transformations/model_test.py
index d3462099..bee469a0 100644
--- a/tests/pipeline_transformations/model_test.py
+++ b/tests/pipeline_transformations/model_test.py
@@ -19,6 +19,7 @@
 from dpsynth.pipeline_transformations import model
 import jax.numpy as jnp
 import mbi
+import numpy as np
 import pandas as pd
 import pipeline_dp
 
@@ -58,7 +59,11 @@ def test_fit_model(self):
   def test_generate_synthetic_data(self, num_records: int | None):
     mock_model = mock.create_autospec(mbi.MarkovRandomField, instance=True)
     mock_model.synthetic_data.return_value = mbi.dataset.Dataset(
-        pd.DataFrame({'col1': [1, 2, 3], 'col2': [4, 5, 6], 'col3': [7, 8, 9]}),
+        data={
+            'col1': np.array([1, 2, 3]),
+            'col2': np.array([4, 5, 6]),
+            'col3': np.array([7, 8, 9]),
+        },
         domain=mbi.Domain(
             attributes=('col2', 'col3', 'col1'), shape=(8, 10, 4)
         ),
diff --git a/tests/pipeline_transformations/mst_test.py b/tests/pipeline_transformations/mst_test.py
index 589dc43b..c231a695 100644
--- a/tests/pipeline_transformations/mst_test.py
+++ b/tests/pipeline_transformations/mst_test.py
@@ -180,7 +180,7 @@ def test_fit_model(self):
     backend = pipeline_dp.LocalBackend()
     num_attributes = 2
     compressed_data = [(0, 1), (0, 1), (1, 0), (1, 1)]
-    compressed_domain = [mbi.Domain(attributes=([0, 1]), shape=(2, 2))]
+    compressed_domain = [mbi.Domain(attributes=(0, 1), shape=(2, 2))]
 
     # Mock one-way DP marginals.
     # For attribute 0: counts for value 0 and 1
diff --git a/tests/pipeline_transformations/numerical_values_derivation_test.py b/tests/pipeline_transformations/numerical_values_derivation_test.py
index b7554d4e..5a75b36f 100644
--- a/tests/pipeline_transformations/numerical_values_derivation_test.py
+++ b/tests/pipeline_transformations/numerical_values_derivation_test.py
@@ -185,6 +185,7 @@ def test_derive_numerical_attributes(self):
     )
 
     accountant.compute_budgets()
+    assert derived_attrs is not None
     derived_attrs_dict = {o.key: o for o in derived_attrs}
 
     self.assertLen(derived_attrs_dict, 2)
@@ -228,6 +229,7 @@ def test_derive_numerical_attributes_empty_input(self):
         num_quantile_buckets=3,
     )
     accountant.compute_budgets()
+    assert derived_attrs is not None
     derived_attrs_list = list(derived_attrs)
     self.assertEmpty(derived_attrs_list)
 
@@ -248,6 +250,7 @@ def test_derive_numerical_attributes_constant_val(self):
     )
 
     accountant.compute_budgets()
+    assert derived_attrs is not None
     derived_attrs_list = list(derived_attrs)
 
     self.assertLen(derived_attrs_list, 1)
diff --git a/tests/relational/domain_test.py b/tests/relational/domain_test.py
index 20aaf3ad..371fc8ee 100644
--- a/tests/relational/domain_test.py
+++ b/tests/relational/domain_test.py
@@ -296,10 +296,22 @@ def test_to_dict_and_roundtrip(self):
 
     # Roundtrip verification
     rt_domains, rt_fks = domain.from_dict(serialized)
+    self.assertIsInstance(
+        rt_domains['households']['income'], base_domain.NumericalAttribute
+    )
+    assert isinstance(
+        rt_domains['households']['income'], base_domain.NumericalAttribute
+    )
     self.assertEqual(
         rt_domains['households']['income'].min_value,
         table_domains['households']['income'].min_value,
     )
+    self.assertIsInstance(
+        rt_domains['persons']['gender'], base_domain.CategoricalAttribute
+    )
+    assert isinstance(
+        rt_domains['persons']['gender'], base_domain.CategoricalAttribute
+    )
     self.assertEqual(
         rt_domains['persons']['gender'].possible_values,
         table_domains['persons']['gender'].possible_values,
@@ -337,6 +349,12 @@ def test_to_yaml_file_and_roundtrip(self):
     rt_domains, rt_fks = domain.from_yaml_file(tmp_path)
     self.assertIn('households', rt_domains)
     self.assertIn('persons', rt_domains)
+    self.assertIsInstance(
+        rt_domains['households']['income'], base_domain.NumericalAttribute
+    )
+    assert isinstance(
+        rt_domains['households']['income'], base_domain.NumericalAttribute
+    )
     self.assertEqual(rt_domains['households']['income'].max_value, 150000.0)
     self.assertEqual(rt_fks, fks)
 
diff --git a/tests/relational/post_processing_test.py b/tests/relational/post_processing_test.py
index 18e46503..38dd8fef 100644
--- a/tests/relational/post_processing_test.py
+++ b/tests/relational/post_processing_test.py
@@ -57,6 +57,7 @@ def test_create_slot_linear_chain_constraints_multi_attribute(self):
     self.assertEqual(c1.domain.shape, (11, 3))
     # Inv. combinations: (10, [0..1]) and ([0..9], 2) -> 2 + 10 = 12 inv. states
     self.assertLen(c1.invalid, 12)
+    assert c1.invalid is not None
     # (10, 0) is mixed state -> must be in invalid
     self.assertTrue(np.any((c1.invalid == [10, 0]).all(axis=1)))
     # (0, 2) is mixed state -> must be in invalid
@@ -424,12 +425,14 @@ def test_unstack_wide_family_records_running_example(self):
         'slot_1.age': 11,
         'slot_2.age': 11,
     })
-    wide_data = {
-        'group_size': np.array([2, 0, 1], dtype=np.int64),
-        'slot_1.age': np.array([5, 10, 3], dtype=np.int64),
-        'slot_2.age': np.array([8, 10, 10], dtype=np.int64),
-    }
-    wide_ds = mbi.Dataset(wide_data, wide_dom)
+    wide_ds = mbi.Dataset(
+        {
+            'group_size': np.array([2, 0, 1], dtype=np.int64),
+            'slot_1.age': np.array([5, 10, 3], dtype=np.int64),
+            'slot_2.age': np.array([8, 10, 10], dtype=np.int64),
+        },
+        wide_dom,
+    )
 
     unstacked_ds, parent_indices = post_processing.unstack_wide_family_records(
         synth_wide_dataset=wide_ds,
@@ -461,12 +464,14 @@ def test_unstack_wide_family_records_size_sliced_strategy(self):
         'slot_1.age': 10,
         'slot_2.age': 10,
     })
-    wide_data = {
-        'group_size': np.array([1, 2], dtype=np.int64),
-        'slot_1.age': np.array([7, 4], dtype=np.int64),
-        'slot_2.age': np.array([7, 9], dtype=np.int64),
-    }
-    wide_ds = mbi.Dataset(wide_data, wide_dom)
+    wide_ds = mbi.Dataset(
+        {
+            'group_size': np.array([1, 2], dtype=np.int64),
+            'slot_1.age': np.array([7, 4], dtype=np.int64),
+            'slot_2.age': np.array([7, 9], dtype=np.int64),
+        },
+        wide_dom,
+    )
 
     unstacked_ds, parent_indices = post_processing.unstack_wide_family_records(
         synth_wide_dataset=wide_ds,
@@ -489,14 +494,16 @@ def test_unstack_wide_family_records_multi_attribute(self):
         'slot_2.age': 11,
         'slot_2.gender': 3,
     })
-    wide_data = {
-        'group_size': np.array([1], dtype=np.int64),
-        'slot_1.age': np.array([4], dtype=np.int64),
-        'slot_1.gender': np.array([1], dtype=np.int64),
-        'slot_2.age': np.array([10], dtype=np.int64),
-        'slot_2.gender': np.array([2], dtype=np.int64),
-    }
-    wide_ds = mbi.Dataset(wide_data, wide_dom)
+    wide_ds = mbi.Dataset(
+        {
+            'group_size': np.array([1], dtype=np.int64),
+            'slot_1.age': np.array([4], dtype=np.int64),
+            'slot_1.gender': np.array([1], dtype=np.int64),
+            'slot_2.age': np.array([10], dtype=np.int64),
+            'slot_2.gender': np.array([2], dtype=np.int64),
+        },
+        wide_dom,
+    )
 
     unstacked_ds, parent_indices = post_processing.unstack_wide_family_records(
         synth_wide_dataset=wide_ds,
@@ -594,11 +601,13 @@ def test_property_slot_linear_chain_constraints_locking_and_treewidth(self):
       for slot in range(1, o + 1):
         for c in constraints:
           if not all(
-              attr.startswith(f'slot_{slot}.') for attr in c.domain.attributes
+              str(attr).startswith(f'slot_{slot}.')
+              for attr in c.domain.attributes
           ):
             continue
           k1, k2 = c.domain.shape[0] - 1, c.domain.shape[1] - 1
           # Monolithic empty: (k1, k2) -> must be VALID (not in invalid)
+          assert c.invalid is not None
           self.assertFalse(np.any((c.invalid == [k1, k2]).all(axis=1)))
           # Monolithic real: (0, 0) -> must be VALID (not in invalid)
           self.assertFalse(np.any((c.invalid == [0, 0]).all(axis=1)))
@@ -658,8 +667,8 @@ def test_property_symmetrize_to_wide_domain_invariants(self):
       pair_copies = [
           m
           for m in expanded
-          if any(a.startswith('slot_') for a in m.clique)
-          and not any(a.startswith('p') for a in m.clique)
+          if any(str(a).startswith('slot_') for a in m.clique)
+          and not any(str(a).startswith('p') for a in m.clique)
       ]
       expected_pair_count = math.comb(s, 2) if (o >= 2 and s >= 2) else 0
       self.assertLen(pair_copies, expected_pair_count)
@@ -727,6 +736,7 @@ def test_property_quantile_copula_coupling_bijection_and_preservation(self):
         )
 
       # 3. Total weight mass strictly preserved
+      assert coupled.weights is not None
       self.assertAlmostEqual(
           float(np.sum(coupled.weights)), float(np.sum(child_weights))
       )
@@ -765,7 +775,7 @@ def test_property_unstack_wide_family_records_consistency(self):
       wide_dom = mbi.Domain(tuple(attrs), tuple(shapes))
 
       group_sizes = rng.integers(0, s + 1, size=n_parents)
-      wide_data = {'group_size': group_sizes}
+      wide_data: dict[str | int, np.ndarray] = {'group_size': group_sizes}
       for slot in range(1, s + 1):
         # Slot active if slot <= group_size
         is_active = slot <= group_sizes
diff --git a/tests/relational/synthesizer_test.py b/tests/relational/synthesizer_test.py
index 359148f0..6fafde11 100644
--- a/tests/relational/synthesizer_test.py
+++ b/tests/relational/synthesizer_test.py
@@ -423,7 +423,7 @@ def test_configure_hyperparameter_validations(self):
       ):
         synthesizer.MultiTableConfig(
             foreign_keys=foreign_keys,
-            exploration_strategy='unsupported_strategy',
+            exploration_strategy='unsupported_strategy',  # pyrefly: ignore[bad-argument-type]
         )
 
     with self.subTest('invalid_discrete_mechanism'):
@@ -930,6 +930,8 @@ def test_encode_and_compress_tables(self):
 
     self.assertIsInstance(datasets['Household'], mbi.Dataset)
     self.assertIsInstance(datasets['Person'], mbi.Dataset)
+    assert datasets['Household'].weights is not None
+    assert datasets['Person'].weights is not None
     np.testing.assert_allclose(datasets['Household'].weights, [1.0, 1.0])
     np.testing.assert_allclose(datasets['Person'].weights, [0.5, 0.5])
 
diff --git a/tests/relational/transformations_test.py b/tests/relational/transformations_test.py
index 99c0ee14..7945e027 100644
--- a/tests/relational/transformations_test.py
+++ b/tests/relational/transformations_test.py
@@ -355,17 +355,18 @@ def test_build_permuted_exploration_dataset_running_example(self):
     # Household 2: (income=30000, region=0), 1 child: P2(age=5)
     # s = 2, o = 2, strategy = 'empty_token'
     parent_dom = mbi.Domain.fromdict({'income': 100000, 'region': 2})
-    parent_data = {
-        'income': np.array([50000, 75000, 30000], dtype=np.int64),
-        'region': np.array([0, 1, 0], dtype=np.int64),
-    }
-    parent_ds = mbi.Dataset(parent_data, parent_dom)
+    parent_ds = mbi.Dataset(
+        {
+            'income': np.array([50000, 75000, 30000], dtype=np.int64),
+            'region': np.array([0, 1, 0], dtype=np.int64),
+        },
+        parent_dom,
+    )
 
     child_dom = mbi.Domain.fromdict({'age': 100})
-    child_data = {
-        'age': np.array([35, 32, 5], dtype=np.int64),
-    }
-    child_ds = mbi.Dataset(child_data, child_dom)
+    child_ds = mbi.Dataset(
+        {'age': np.array([35, 32, 5], dtype=np.int64)}, child_dom
+    )
 
     parent_pks = ['H0', 'H1', 'H2']
     child_fks = ['H0', 'H0', 'H2']
@@ -388,6 +389,7 @@ def test_build_permuted_exploration_dataset_running_example(self):
     self.assertEqual(expl_ds.records, 5)
 
     # 1. Weight mass invariant: sum(weights) == N_parents = 3.0
+    assert expl_ds.weights is not None
     self.assertAlmostEqual(float(np.sum(expl_ds.weights)), 3.0)
 
     # 2. Domain check
@@ -435,6 +437,7 @@ def test_build_permuted_exploration_dataset_size_sliced(self):
         strategy='size_sliced',
     )
     self.assertEqual(expl_ds.records, 2)
+    assert expl_ds.weights is not None
     self.assertAlmostEqual(float(np.sum(expl_ds.weights)), 2.0)
     self.assertEqual(expl_ds.domain.shape, (100, 3, 100, 100))
 
@@ -449,6 +452,7 @@ def test_build_permuted_exploration_dataset_empty_and_orphans(self):
         parent_empty, child_empty, [], [], max_group_size=2
     )
     self.assertEqual(expl_empty.records, 0)
+    assert expl_empty.weights is not None
     self.assertAlmostEqual(float(np.sum(expl_empty.weights)), 0.0)
 
     # All orphan children
diff --git a/tests/text/bulk_inference_test.py b/tests/text/bulk_inference_test.py
index 3c0bba6e..9d5dbfa3 100644
--- a/tests/text/bulk_inference_test.py
+++ b/tests/text/bulk_inference_test.py
@@ -417,7 +417,9 @@ def test_categorical_with_description(self):
 
   def test_categorical_rejects_none(self):
     with self.assertRaises(ValueError):
-      domain.CategoricalAttribute(possible_values=[None, 'A', 'B'])
+      domain.CategoricalAttribute(
+          possible_values=[None, 'A', 'B']  # pyrefly: ignore[bad-argument-type]
+      )
 
   def test_numerical_produces_number(self):
     d = {
diff --git a/tests/text/dp_sft_test.py b/tests/text/dp_sft_test.py
index 6c7eceb6..a9246bf6 100644
--- a/tests/text/dp_sft_test.py
+++ b/tests/text/dp_sft_test.py
@@ -132,7 +132,9 @@ def test_custom_checkpoint_path(self):
 
   def test_unknown_name_raises(self):
     with self.assertRaises(ValueError):
-      model.GemmaModel.default('nonexistent')
+      model.GemmaModel.default(
+          'nonexistent'  # pyrefly: ignore[bad-argument-type]
+      )
 
 
 class LoraConfigTest(absltest.TestCase):
diff --git a/tests/text/dp_trainer_test.py b/tests/text/dp_trainer_test.py
index 0be08e2e..a73b196c 100644
--- a/tests/text/dp_trainer_test.py
+++ b/tests/text/dp_trainer_test.py
@@ -26,9 +26,9 @@ def _dummy_params_and_loss():
   """Creates a trivial pytree and loss function for testing."""
   params = {'w': jnp.ones((4, 4))}
 
-  def loss_fn(params, batch, prng):
+  def loss_fn(params, data, prng):
     del prng
-    return jnp.sum(params['w'] * batch['x']), ()
+    return jnp.sum(params['w'] * data['x']), ()
 
   return params, loss_fn
 
@@ -52,6 +52,12 @@ def test_default_creates_valid_config(self):
         mechanism_config=self.dpsgd_config,
         optimizer=optax.adamw(1e-4),
     )
+    self.assertIsInstance(
+        trainer.mechanism_config, jax_privacy.execution_plan.BandMFConfig
+    )
+    assert isinstance(
+        trainer.mechanism_config, jax_privacy.execution_plan.BandMFConfig
+    )
     self.assertEqual(trainer.mechanism_config.iterations, 100)
     self.assertIsNone(trainer.mechanism_config.noise_multiplier)
 
@@ -64,6 +70,14 @@ def test_calibrate_sets_noise_multiplier(self):
         optimizer=optax.adamw(1e-4),
     ).configure(zcdp_rho=0.5)
     # Single band: sigma = sqrt(T / (2 * rho)) = 10.
+    self.assertIsInstance(
+        trainer.mechanism_config, jax_privacy.execution_plan.BandMFConfig
+    )
+    assert isinstance(
+        trainer.mechanism_config, jax_privacy.execution_plan.BandMFConfig
+    )
+    self.assertIsNotNone(trainer.mechanism_config.noise_multiplier)
+    assert trainer.mechanism_config.noise_multiplier is not None
     self.assertAlmostEqual(trainer.mechanism_config.noise_multiplier, 10.0)
     self.assertIsNotNone(trainer.dp_event)
 
@@ -88,10 +102,10 @@ def test_nnx_split_merge_round_trip(self):
 
     graphdef, trainable, frozen = nnx.split(lora_model, nnx.LoRAParam, ...)
 
-    def loss_fn(params, batch, prng):
+    def loss_fn(params, data, prng):
       del prng
       model = nnx.merge(graphdef, params, frozen)
-      x = batch['x']
+      x = data['x']
       return jnp.mean(model(x) ** 2), ()
 
     full_batch_config = jax_privacy.execution_plan.BandMFConfig.default(
@@ -101,7 +115,7 @@ def loss_fn(params, batch, prng):
         l2_clip_norm=1.0,
     )
     trainer = dp_trainer.DPTrainer(
-        init_params=trainable,
+        init_params=trainable,  # pyrefly: ignore[bad-argument-type]
         loss_fn=loss_fn,
         mechanism_config=full_batch_config,
         optimizer=optax.adamw(1e-4),