alberta-framework 0.2.1__tar.gz → 0.3.0__tar.gz

{alberta_framework-0.2.1 → alberta_framework-0.3.0}/.github/workflows/publish.yml

@@ -4,6 +4,7 @@ on:
   push:
     tags:
       - "v*"
+  workflow_dispatch:
 
 permissions:
   id-token: write

{alberta_framework-0.2.1 → alberta_framework-0.3.0}/CLAUDE.md

@@ -70,7 +70,8 @@ mkdocs build          # Build static site to site/
 ## Development Guidelines
 
 ### Design Principles
-- **Immutable State**: All state uses NamedTuples for JAX compatibility
+- **Immutable State**: All state uses `@chex.dataclass(frozen=True)` for JAX PyTree compatibility
+- **Type Safety**: jaxtyping annotations for shape checking (`Float[Array, " feature_dim"]`)
 - **Functional Style**: Pure functions enable `jit`, `vmap`, `jax.lax.scan`
 - **Scan-Based Learning**: Learning loops use `jax.lax.scan` for JIT-compiled training
 - **Composition**: Learners accept optimizers as parameters
@@ -85,6 +86,7 @@ mkdocs build          # Build static site to site/
 ### Testing
 - Tests are in `tests/` directory
 - Use pytest fixtures from `conftest.py`
+- Use chex assertions: `chex.assert_shape()`, `chex.assert_trees_all_close()`, `chex.assert_tree_all_finite()`
 - All tests should pass before committing
 
 ## Key Algorithms
@@ -465,3 +467,16 @@ The publish workflow uses OpenID Connect (no API tokens). Configure on PyPI:
 1. PyPI project → Settings → Publishing → Add GitHub publisher
 2. Repository: `j-klawson/alberta-framework`, Workflow: `publish.yml`, Environment: `pypi`
 3. Repeat on TestPyPI with environment: `testpypi`
+
+## Changelog
+
+### v0.3.0 (2026-02-03)
+- **FEATURE**: Migrated all state types from NamedTuple to `@chex.dataclass(frozen=True)` for DeepMind-style JAX compatibility
+- **FEATURE**: Added jaxtyping shape annotations for compile-time type safety (`Float[Array, " feature_dim"]`, `PRNGKeyArray`, etc.)
+- **FEATURE**: Updated test suite to use chex assertions (`chex.assert_shape`, `chex.assert_tree_all_finite`, `chex.assert_trees_all_close`)
+- **DEPS**: Added `chex>=0.1.86` and `jaxtyping>=0.2.28` as required dependencies
+- **DEPS**: Added `beartype>=0.18.0` as optional dev dependency for runtime type checking
+
+### v0.2.2 (2026-02-02)
+- Fixed mypy type errors in `run_learning_loop_batched` and `run_normalized_learning_loop_batched` functions
+- Added `typing.cast` to properly handle conditional return type unpacking in batched learning loops

{alberta_framework-0.2.1 → alberta_framework-0.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: alberta-framework
-Version: 0.2.1
+Version: 0.3.0
 Summary: Implementation of the Alberta Plan for AI Research - continual learning with meta-learned step-sizes
 Project-URL: Homepage, https://github.com/j-klawson/alberta-framework
 Project-URL: Repository, https://github.com/j-klawson/alberta-framework
@@ -17,14 +17,17 @@ Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.13
 Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Requires-Python: >=3.13
+Requires-Dist: chex>=0.1.86
 Requires-Dist: jax>=0.4
 Requires-Dist: jaxlib>=0.4
+Requires-Dist: jaxtyping>=0.2.28
 Requires-Dist: joblib>=1.3
 Requires-Dist: matplotlib>=3.8
 Requires-Dist: numpy>=2.0
 Requires-Dist: scipy>=1.11
 Requires-Dist: tqdm>=4.66
 Provides-Extra: dev
+Requires-Dist: beartype>=0.18.0; extra == 'dev'
 Requires-Dist: mypy; extra == 'dev'
 Requires-Dist: pytest-cov; extra == 'dev'
 Requires-Dist: pytest>=8.0; extra == 'dev'

{alberta_framework-0.2.1 → alberta_framework-0.3.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "alberta-framework"
-version = "0.2.1"
+version = "0.3.0"
 description = "Implementation of the Alberta Plan for AI Research - continual learning with meta-learned step-sizes"
 readme = "README.md"
 license = "Apache-2.0"
@@ -29,6 +29,8 @@ dependencies = [
     "scipy>=1.11",
     "joblib>=1.3",
     "tqdm>=4.66",
+    "chex>=0.1.86",
+    "jaxtyping>=0.2.28",
 ]
 
 [project.optional-dependencies]
@@ -37,6 +39,7 @@ dev = [
     "pytest-cov",
     "ruff",
     "mypy",
+    "beartype>=0.18.0",
 ]
 docs = [
     "mkdocs>=1.6",
@@ -99,6 +102,39 @@ ignore_missing_imports = true
 module = "pandas.*"
 ignore_missing_imports = true
 
+[[tool.mypy.overrides]]
+module = "chex.*"
+ignore_missing_imports = true
+
+[[tool.mypy.overrides]]
+module = "jaxtyping.*"
+ignore_missing_imports = true
+
+# chex dataclasses don't have full mypy support - disable strict checks for core modules
+[[tool.mypy.overrides]]
+module = "alberta_framework.core.types"
+disable_error_code = ["call-arg"]
+
+[[tool.mypy.overrides]]
+module = "alberta_framework.core.normalizers"
+disable_error_code = ["call-arg"]
+
+[[tool.mypy.overrides]]
+module = "alberta_framework.core.optimizers"
+disable_error_code = ["call-arg"]
+
+[[tool.mypy.overrides]]
+module = "alberta_framework.core.learners"
+disable_error_code = ["call-arg", "arg-type", "union-attr"]
+
+[[tool.mypy.overrides]]
+module = "alberta_framework.streams.synthetic"
+disable_error_code = ["call-arg"]
+
+[[tool.mypy.overrides]]
+module = "alberta_framework.streams.gymnasium"
+disable_error_code = ["call-arg"]
+
 [tool.pytest.ini_options]
 testpaths = ["tests"]
 addopts = "-v"

{alberta_framework-0.2.1 → alberta_framework-0.3.0}/src/alberta_framework/core/learners.py

@@ -5,11 +5,13 @@ for temporally-uniform learning. Uses JAX's scan for efficient JIT-compiled
 training loops.
 """
 
-from typing import NamedTuple
+from typing import cast
 
+import chex
 import jax
 import jax.numpy as jnp
 from jax import Array
+from jaxtyping import Float
 
 from alberta_framework.core.normalizers import NormalizerState, OnlineNormalizer
 from alberta_framework.core.optimizers import LMS, Optimizer
@@ -33,7 +35,9 @@ from alberta_framework.streams.base import ScanStream
 # Type alias for any optimizer type
 AnyOptimizer = Optimizer[LMSState] | Optimizer[IDBDState] | Optimizer[AutostepState]
 
-class UpdateResult(NamedTuple):
+
+@chex.dataclass(frozen=True)
+class UpdateResult:
     """Result of a learner update step.
 
     Attributes:
@@ -45,8 +49,38 @@ class UpdateResult(NamedTuple):
 
     state: LearnerState
     prediction: Prediction
-    error: Array
-    metrics: Array
+    error: Float[Array, ""]
+    metrics: Float[Array, " 3"]
+
+
+@chex.dataclass(frozen=True)
+class NormalizedLearnerState:
+    """State for a learner with online feature normalization.
+
+    Attributes:
+        learner_state: Underlying learner state (weights, bias, optimizer)
+        normalizer_state: Online normalizer state (mean, var estimates)
+    """
+
+    learner_state: LearnerState
+    normalizer_state: NormalizerState
+
+
+@chex.dataclass(frozen=True)
+class NormalizedUpdateResult:
+    """Result of a normalized learner update step.
+
+    Attributes:
+        state: Updated normalized learner state
+        prediction: Prediction made before update
+        error: Prediction error
+        metrics: Array of metrics [squared_error, error, mean_step_size, normalizer_mean_var]
+    """
+
+    state: NormalizedLearnerState
+    prediction: Prediction
+    error: Float[Array, ""]
+    metrics: Float[Array, " 4"]
 
 
 class LinearLearner:
@@ -133,8 +167,7 @@ class LinearLearner:
         # Note: type ignore needed because we can't statically prove optimizer_state
         # matches the optimizer's expected state type (though they will at runtime)
         opt_update = self._optimizer.update(
-            state.optimizer_state,  # type: ignore[arg-type]
-            error,
+            state.optimizer_state,            error,
             observation,
         )
 
@@ -275,12 +308,12 @@ def run_learning_loop[StreamStateT](
             opt_state = result.state.optimizer_state
             if hasattr(opt_state, "log_step_sizes"):
                 # IDBD stores log step-sizes
-                weight_ss = jnp.exp(opt_state.log_step_sizes)  # type: ignore[union-attr]
-                bias_ss = opt_state.bias_step_size  # type: ignore[union-attr]
+                weight_ss = jnp.exp(opt_state.log_step_sizes)
+                bias_ss = opt_state.bias_step_size
             elif hasattr(opt_state, "step_sizes"):
                 # Autostep stores step-sizes directly
-                weight_ss = opt_state.step_sizes  # type: ignore[union-attr]
-                bias_ss = opt_state.bias_step_size  # type: ignore[union-attr]
+                weight_ss = opt_state.step_sizes
+                bias_ss = opt_state.bias_step_size
             else:
                 # LMS has a single fixed step-size
                 weight_ss = jnp.full(feature_dim, opt_state.step_size)
@@ -316,8 +349,7 @@ def run_learning_loop[StreamStateT](
                 new_norm_history = jax.lax.cond(
                     should_record,
                     lambda _: norm_history.at[recording_idx].set(
-                        opt_state.normalizers  # type: ignore[union-attr]
-                    ),
+                        opt_state.normalizers                    ),
                     lambda _: norm_history,
                     None,
                 )
@@ -361,34 +393,6 @@ def run_learning_loop[StreamStateT](
         return final_learner, metrics, history
 
 
-class NormalizedLearnerState(NamedTuple):
-    """State for a learner with online feature normalization.
-
-    Attributes:
-        learner_state: Underlying learner state (weights, bias, optimizer)
-        normalizer_state: Online normalizer state (mean, var estimates)
-    """
-
-    learner_state: LearnerState
-    normalizer_state: NormalizerState
-
-
-class NormalizedUpdateResult(NamedTuple):
-    """Result of a normalized learner update step.
-
-    Attributes:
-        state: Updated normalized learner state
-        prediction: Prediction made before update
-        error: Prediction error
-        metrics: Array of metrics [squared_error, error, mean_step_size, normalizer_mean_var]
-    """
-
-    state: NormalizedLearnerState
-    prediction: Prediction
-    error: Array
-    metrics: Array
-
-
 class NormalizedLinearLearner:
     """Linear learner with online feature normalization.
 
@@ -702,12 +706,12 @@ def run_normalized_learning_loop[StreamStateT](
             opt_state = result.state.learner_state.optimizer_state
             if hasattr(opt_state, "log_step_sizes"):
                 # IDBD stores log step-sizes
-                weight_ss = jnp.exp(opt_state.log_step_sizes)  # type: ignore[union-attr]
-                bias_ss = opt_state.bias_step_size  # type: ignore[union-attr]
+                weight_ss = jnp.exp(opt_state.log_step_sizes)
+                bias_ss = opt_state.bias_step_size
             elif hasattr(opt_state, "step_sizes"):
                 # Autostep stores step-sizes directly
-                weight_ss = opt_state.step_sizes  # type: ignore[union-attr]
-                bias_ss = opt_state.bias_step_size  # type: ignore[union-attr]
+                weight_ss = opt_state.step_sizes
+                bias_ss = opt_state.bias_step_size
             else:
                 # LMS has a single fixed step-size
                 weight_ss = jnp.full(feature_dim, opt_state.step_size)
@@ -741,8 +745,7 @@ def run_normalized_learning_loop[StreamStateT](
                 new_ss_norm = jax.lax.cond(
                     should_record_ss,
                     lambda _: ss_norm.at[recording_idx].set(
-                        opt_state.normalizers  # type: ignore[union-attr]
-                    ),
+                        opt_state.normalizers                    ),
                     lambda _: ss_norm,
                     None,
                 )
@@ -825,17 +828,14 @@ def run_normalized_learning_loop[StreamStateT](
         ss_history_result = StepSizeHistory(
             step_sizes=final_ss_hist,
             bias_step_sizes=final_ss_bias_hist,
-            recording_indices=final_ss_rec,  # type: ignore[arg-type]
-            normalizers=final_ss_norm,
+            recording_indices=final_ss_rec,            normalizers=final_ss_norm,
         )
 
     norm_history_result = None
     if normalizer_tracking is not None and final_n_means is not None:
         norm_history_result = NormalizerHistory(
             means=final_n_means,
-            variances=final_n_vars,  # type: ignore[arg-type]
-            recording_indices=final_n_rec,  # type: ignore[arg-type]
-        )
+            variances=final_n_vars,            recording_indices=final_n_rec,        )
 
     # Return appropriate tuple based on what was tracked
     if ss_history_result is not None and norm_history_result is not None:
@@ -900,10 +900,12 @@ def run_learning_loop_batched[StreamStateT](
             learner, stream, num_steps, key, learner_state, step_size_tracking
         )
         if step_size_tracking is not None:
-            state, metrics, history = result
+            state, metrics, history = cast(
+                tuple[LearnerState, Array, StepSizeHistory], result
+            )
             return state, metrics, history
         else:
-            state, metrics = result
+            state, metrics = cast(tuple[LearnerState, Array], result)
             # Return None for history to maintain consistent output structure
             return state, metrics, None
 
@@ -911,7 +913,7 @@ def run_learning_loop_batched[StreamStateT](
     batched_states, batched_metrics, batched_history = jax.vmap(single_seed_run)(keys)
 
     # Reconstruct batched history if tracking was enabled
-    if step_size_tracking is not None:
+    if step_size_tracking is not None and batched_history is not None:
         batched_step_size_history = StepSizeHistory(
             step_sizes=batched_history.step_sizes,
             bias_step_sizes=batched_history.bias_step_sizes,
@@ -993,16 +995,23 @@ def run_normalized_learning_loop_batched[StreamStateT](
 
         # Unpack based on what tracking was enabled
         if step_size_tracking is not None and normalizer_tracking is not None:
-            state, metrics, ss_history, norm_history = result
+            state, metrics, ss_history, norm_history = cast(
+                tuple[NormalizedLearnerState, Array, StepSizeHistory, NormalizerHistory],
+                result,
+            )
             return state, metrics, ss_history, norm_history
         elif step_size_tracking is not None:
-            state, metrics, ss_history = result
+            state, metrics, ss_history = cast(
+                tuple[NormalizedLearnerState, Array, StepSizeHistory], result
+            )
             return state, metrics, ss_history, None
         elif normalizer_tracking is not None:
-            state, metrics, norm_history = result
+            state, metrics, norm_history = cast(
+                tuple[NormalizedLearnerState, Array, NormalizerHistory], result
+            )
             return state, metrics, None, norm_history
         else:
-            state, metrics = result
+            state, metrics = cast(tuple[NormalizedLearnerState, Array], result)
             return state, metrics, None, None
 
     # vmap over the keys dimension
@@ -1011,7 +1020,7 @@ def run_normalized_learning_loop_batched[StreamStateT](
     )
 
     # Reconstruct batched histories if tracking was enabled
-    if step_size_tracking is not None:
+    if step_size_tracking is not None and batched_ss_history is not None:
         batched_step_size_history = StepSizeHistory(
             step_sizes=batched_ss_history.step_sizes,
             bias_step_sizes=batched_ss_history.bias_step_sizes,
@@ -1021,7 +1030,7 @@ def run_normalized_learning_loop_batched[StreamStateT](
     else:
         batched_step_size_history = None
 
-    if normalizer_tracking is not None:
+    if normalizer_tracking is not None and batched_norm_history is not None:
         batched_normalizer_history = NormalizerHistory(
             means=batched_norm_history.means,
             variances=batched_norm_history.variances,

{alberta_framework-0.2.1 → alberta_framework-0.3.0}/src/alberta_framework/core/normalizers.py

@@ -6,13 +6,14 @@ and variance at every time step, following the principle of temporal uniformity.
 Reference: Welford's online algorithm for numerical stability.
 """
 
-from typing import NamedTuple
-
+import chex
 import jax.numpy as jnp
 from jax import Array
+from jaxtyping import Float
 
 
-class NormalizerState(NamedTuple):
+@chex.dataclass(frozen=True)
+class NormalizerState:
     """State for online feature normalization.
 
     Uses Welford's online algorithm for numerically stable estimation
@@ -25,10 +26,10 @@ class NormalizerState(NamedTuple):
         decay: Exponential decay factor for estimates (1.0 = no decay, pure online)
     """
 
-    mean: Array  # Shape: (feature_dim,)
-    var: Array  # Shape: (feature_dim,)
-    sample_count: Array  # Scalar
-    decay: Array  # Scalar
+    mean: Float[Array, " feature_dim"]
+    var: Float[Array, " feature_dim"]
+    sample_count: Float[Array, ""]
+    decay: Float[Array, ""]
 
 
 class OnlineNormalizer:

{alberta_framework-0.2.1 → alberta_framework-0.3.0}/src/alberta_framework/core/optimizers.py

@@ -10,15 +10,17 @@ References:
 """
 
 from abc import ABC, abstractmethod
-from typing import NamedTuple
 
+import chex
 import jax.numpy as jnp
 from jax import Array
+from jaxtyping import Float
 
 from alberta_framework.core.types import AutostepState, IDBDState, LMSState
 
 
-class OptimizerUpdate(NamedTuple):
+@chex.dataclass(frozen=True)
+class OptimizerUpdate:
     """Result of an optimizer update step.
 
     Attributes:
@@ -28,8 +30,8 @@ class OptimizerUpdate(NamedTuple):
         metrics: Dictionary of metrics for logging (values are JAX arrays for scan compatibility)
     """
 
-    weight_delta: Array
-    bias_delta: Array
+    weight_delta: Float[Array, " feature_dim"]
+    bias_delta: Float[Array, ""]
     new_state: LMSState | IDBDState | AutostepState
     metrics: dict[str, Array]
 

{alberta_framework-0.2.1 → alberta_framework-0.3.0}/src/alberta_framework/core/types.py

@@ -1,13 +1,15 @@
 """Type definitions for the Alberta Framework.
 
 This module defines the core data types used throughout the framework,
-following JAX conventions with immutable NamedTuples for state management.
+using chex dataclasses for JAX compatibility and jaxtyping for shape annotations.
 """
 
-from typing import TYPE_CHECKING, NamedTuple
+from typing import TYPE_CHECKING
 
+import chex
 import jax.numpy as jnp
 from jax import Array
+from jaxtyping import Float, Int, PRNGKeyArray
 
 if TYPE_CHECKING:
     from alberta_framework.core.learners import NormalizedLearnerState
@@ -19,7 +21,8 @@ Prediction = Array  # y_t: model output
 Reward = float  # r_t: scalar reward
 
 
-class TimeStep(NamedTuple):
+@chex.dataclass(frozen=True)
+class TimeStep:
     """Single experience from an experience stream.
 
     Attributes:
@@ -27,25 +30,12 @@ class TimeStep(NamedTuple):
         target: Desired output y*_t (for supervised learning)
     """
 
-    observation: Observation
-    target: Target
+    observation: Float[Array, " feature_dim"]
+    target: Float[Array, " 1"]
 
 
-class LearnerState(NamedTuple):
-    """State for a linear learner.
-
-    Attributes:
-        weights: Weight vector for linear prediction
-        bias: Bias term
-        optimizer_state: State maintained by the optimizer
-    """
-
-    weights: Array
-    bias: Array
-    optimizer_state: "LMSState | IDBDState | AutostepState"
-
-
-class LMSState(NamedTuple):
+@chex.dataclass(frozen=True)
+class LMSState:
     """State for the LMS (Least Mean Square) optimizer.
 
     LMS uses a fixed step-size, so state only tracks the step-size parameter.
@@ -54,10 +44,11 @@ class LMSState(NamedTuple):
         step_size: Fixed learning rate alpha
     """
 
-    step_size: Array
+    step_size: Float[Array, ""]
 
 
-class IDBDState(NamedTuple):
+@chex.dataclass(frozen=True)
+class IDBDState:
     """State for the IDBD (Incremental Delta-Bar-Delta) optimizer.
 
     IDBD maintains per-weight adaptive step-sizes that are meta-learned
@@ -73,14 +64,15 @@ class IDBDState(NamedTuple):
         bias_trace: Trace for the bias term
     """
 
-    log_step_sizes: Array  # log(alpha_i) for numerical stability
-    traces: Array  # h_i: trace of weight-feature products
-    meta_step_size: Array  # beta: step-size for the step-sizes
-    bias_step_size: Array  # Step-size for bias
-    bias_trace: Array  # Trace for bias
+    log_step_sizes: Float[Array, " feature_dim"]  # log(alpha_i) for numerical stability
+    traces: Float[Array, " feature_dim"]  # h_i: trace of weight-feature products
+    meta_step_size: Float[Array, ""]  # beta: step-size for the step-sizes
+    bias_step_size: Float[Array, ""]  # Step-size for bias
+    bias_trace: Float[Array, ""]  # Trace for bias
 
 
-class AutostepState(NamedTuple):
+@chex.dataclass(frozen=True)
+class AutostepState:
     """State for the Autostep optimizer.
 
     Autostep is a tuning-free step-size adaptation algorithm that normalizes
@@ -100,17 +92,33 @@ class AutostepState(NamedTuple):
         bias_normalizer: Normalizer for the bias gradient
     """
 
-    step_sizes: Array  # alpha_i
-    traces: Array  # h_i
-    normalizers: Array  # v_i: running max of |gradient|
-    meta_step_size: Array  # mu
-    normalizer_decay: Array  # tau
-    bias_step_size: Array
-    bias_trace: Array
-    bias_normalizer: Array
+    step_sizes: Float[Array, " feature_dim"]  # alpha_i
+    traces: Float[Array, " feature_dim"]  # h_i
+    normalizers: Float[Array, " feature_dim"]  # v_i: running max of |gradient|
+    meta_step_size: Float[Array, ""]  # mu
+    normalizer_decay: Float[Array, ""]  # tau
+    bias_step_size: Float[Array, ""]
+    bias_trace: Float[Array, ""]
+    bias_normalizer: Float[Array, ""]
+
+
+@chex.dataclass(frozen=True)
+class LearnerState:
+    """State for a linear learner.
+
+    Attributes:
+        weights: Weight vector for linear prediction
+        bias: Bias term
+        optimizer_state: State maintained by the optimizer
+    """
+
+    weights: Float[Array, " feature_dim"]
+    bias: Float[Array, ""]
+    optimizer_state: LMSState | IDBDState | AutostepState
 
 
-class StepSizeTrackingConfig(NamedTuple):
+@chex.dataclass(frozen=True)
+class StepSizeTrackingConfig:
     """Configuration for recording per-weight step-sizes during training.
 
     Attributes:
@@ -122,7 +130,8 @@ class StepSizeTrackingConfig(NamedTuple):
     include_bias: bool = True
 
 
-class StepSizeHistory(NamedTuple):
+@chex.dataclass(frozen=True)
+class StepSizeHistory:
     """History of per-weight step-sizes recorded during training.
 
     Attributes:
@@ -133,13 +142,14 @@ class StepSizeHistory(NamedTuple):
             shape (num_recordings, num_weights) or None. Only populated for Autostep optimizer.
     """
 
-    step_sizes: Array  # (num_recordings, num_weights)
-    bias_step_sizes: Array | None  # (num_recordings,) or None
-    recording_indices: Array  # (num_recordings,)
-    normalizers: Array | None = None  # (num_recordings, num_weights) - Autostep v_i
+    step_sizes: Float[Array, "num_recordings feature_dim"]
+    bias_step_sizes: Float[Array, " num_recordings"] | None
+    recording_indices: Int[Array, " num_recordings"]
+    normalizers: Float[Array, "num_recordings feature_dim"] | None = None
 
 
-class NormalizerTrackingConfig(NamedTuple):
+@chex.dataclass(frozen=True)
+class NormalizerTrackingConfig:
     """Configuration for recording per-feature normalizer state during training.
 
     Attributes:
@@ -149,7 +159,8 @@ class NormalizerTrackingConfig(NamedTuple):
     interval: int
 
 
-class NormalizerHistory(NamedTuple):
+@chex.dataclass(frozen=True)
+class NormalizerHistory:
     """History of per-feature normalizer state recorded during training.
 
     Used for analyzing how the OnlineNormalizer adapts to distribution shifts
@@ -162,12 +173,13 @@ class NormalizerHistory(NamedTuple):
         recording_indices: Step indices where recordings were made, shape (num_recordings,)
     """
 
-    means: Array  # (num_recordings, feature_dim)
-    variances: Array  # (num_recordings, feature_dim)
-    recording_indices: Array  # (num_recordings,)
+    means: Float[Array, "num_recordings feature_dim"]
+    variances: Float[Array, "num_recordings feature_dim"]
+    recording_indices: Int[Array, " num_recordings"]
 
 
-class BatchedLearningResult(NamedTuple):
+@chex.dataclass(frozen=True)
+class BatchedLearningResult:
     """Result from batched learning loop across multiple seeds.
 
     Used with `run_learning_loop_batched` for vmap-based GPU parallelization.
@@ -180,12 +192,13 @@ class BatchedLearningResult(NamedTuple):
             or None if tracking was disabled
     """
 
-    states: "LearnerState"  # Batched: each array has shape (num_seeds, ...)
-    metrics: Array  # Shape: (num_seeds, num_steps, 3)
+    states: LearnerState  # Batched: each array has shape (num_seeds, ...)
+    metrics: Float[Array, "num_seeds num_steps 3"]
     step_size_history: StepSizeHistory | None
 
 
-class BatchedNormalizedResult(NamedTuple):
+@chex.dataclass(frozen=True)
+class BatchedNormalizedResult:
     """Result from batched normalized learning loop across multiple seeds.
 
     Used with `run_normalized_learning_loop_batched` for vmap-based GPU parallelization.
@@ -201,7 +214,7 @@ class BatchedNormalizedResult(NamedTuple):
     """
 
     states: "NormalizedLearnerState"  # Batched: each array has shape (num_seeds, ...)
-    metrics: Array  # Shape: (num_seeds, num_steps, 4)
+    metrics: Float[Array, "num_seeds num_steps 4"]
     step_size_history: StepSizeHistory | None
     normalizer_history: NormalizerHistory | None