alberta-framework 0.2.1__py3-none-any.whl → 0.3.0__py3-none-any.whl

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/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/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/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
 

alberta_framework/streams/synthetic.py

@@ -7,17 +7,20 @@ track and adapt.
 All streams use JAX-compatible pure functions that work with jax.lax.scan.
 """
 
-from typing import Any, NamedTuple
+from typing import Any
 
+import chex
 import jax.numpy as jnp
 import jax.random as jr
 from jax import Array
+from jaxtyping import Float, Int, PRNGKeyArray
 
 from alberta_framework.core.types import TimeStep
 from alberta_framework.streams.base import ScanStream
 
 
-class RandomWalkState(NamedTuple):
+@chex.dataclass(frozen=True)
+class RandomWalkState:
     """State for RandomWalkStream.
 
     Attributes:
@@ -25,8 +28,8 @@ class RandomWalkState(NamedTuple):
         true_weights: Current true target weights
     """
 
-    key: Array
-    true_weights: Array
+    key: PRNGKeyArray
+    true_weights: Float[Array, " feature_dim"]
 
 
 class RandomWalkStream:
@@ -110,7 +113,8 @@ class RandomWalkStream:
         return timestep, new_state
 
 
-class AbruptChangeState(NamedTuple):
+@chex.dataclass(frozen=True)
+class AbruptChangeState:
     """State for AbruptChangeStream.
 
     Attributes:
@@ -119,9 +123,9 @@ class AbruptChangeState(NamedTuple):
         step_count: Number of steps taken
     """
 
-    key: Array
-    true_weights: Array
-    step_count: Array
+    key: PRNGKeyArray
+    true_weights: Float[Array, " feature_dim"]
+    step_count: Int[Array, ""]
 
 
 class AbruptChangeStream:
@@ -219,7 +223,8 @@ class AbruptChangeStream:
         return timestep, new_state
 
 
-class SuttonExperiment1State(NamedTuple):
+@chex.dataclass(frozen=True)
+class SuttonExperiment1State:
     """State for SuttonExperiment1Stream.
 
     Attributes:
@@ -228,9 +233,9 @@ class SuttonExperiment1State(NamedTuple):
         step_count: Number of steps taken
     """
 
-    key: Array
-    signs: Array
-    step_count: Array
+    key: PRNGKeyArray
+    signs: Float[Array, " num_relevant"]
+    step_count: Int[Array, ""]
 
 
 class SuttonExperiment1Stream:
@@ -341,7 +346,8 @@ class SuttonExperiment1Stream:
         return timestep, new_state
 
 
-class CyclicState(NamedTuple):
+@chex.dataclass(frozen=True)
+class CyclicState:
     """State for CyclicStream.
 
     Attributes:
@@ -350,9 +356,9 @@ class CyclicState(NamedTuple):
         step_count: Number of steps taken
     """
 
-    key: Array
-    configurations: Array
-    step_count: Array
+    key: PRNGKeyArray
+    configurations: Float[Array, "num_configs feature_dim"]
+    step_count: Int[Array, ""]
 
 
 class CyclicStream:
@@ -452,7 +458,8 @@ class CyclicStream:
         return timestep, new_state
 
 
-class PeriodicChangeState(NamedTuple):
+@chex.dataclass(frozen=True)
+class PeriodicChangeState:
     """State for PeriodicChangeStream.
 
     Attributes:
@@ -462,10 +469,10 @@ class PeriodicChangeState(NamedTuple):
         step_count: Number of steps taken
     """
 
-    key: Array
-    base_weights: Array
-    phases: Array
-    step_count: Array
+    key: PRNGKeyArray
+    base_weights: Float[Array, " feature_dim"]
+    phases: Float[Array, " feature_dim"]
+    step_count: Int[Array, ""]
 
 
 class PeriodicChangeStream:
@@ -573,7 +580,8 @@ class PeriodicChangeStream:
         return timestep, new_state
 
 
-class ScaledStreamState(NamedTuple):
+@chex.dataclass(frozen=True)
+class ScaledStreamState:
     """State for ScaledStreamWrapper.
 
     Attributes:
@@ -714,7 +722,8 @@ def make_scale_range(
         return jnp.linspace(min_scale, max_scale, feature_dim, dtype=jnp.float32)
 
 
-class DynamicScaleShiftState(NamedTuple):
+@chex.dataclass(frozen=True)
+class DynamicScaleShiftState:
     """State for DynamicScaleShiftStream.
 
     Attributes:
@@ -724,10 +733,10 @@ class DynamicScaleShiftState(NamedTuple):
         step_count: Number of steps taken
     """
 
-    key: Array
-    true_weights: Array
-    current_scales: Array
-    step_count: Array
+    key: PRNGKeyArray
+    true_weights: Float[Array, " feature_dim"]
+    current_scales: Float[Array, " feature_dim"]
+    step_count: Int[Array, ""]
 
 
 class DynamicScaleShiftStream:
@@ -858,7 +867,8 @@ class DynamicScaleShiftStream:
         return timestep, new_state
 
 
-class ScaleDriftState(NamedTuple):
+@chex.dataclass(frozen=True)
+class ScaleDriftState:
     """State for ScaleDriftStream.
 
     Attributes:
@@ -868,10 +878,10 @@ class ScaleDriftState(NamedTuple):
         step_count: Number of steps taken
     """
 
-    key: Array
-    true_weights: Array
-    log_scales: Array
-    step_count: Array
+    key: PRNGKeyArray
+    true_weights: Float[Array, " feature_dim"]
+    log_scales: Float[Array, " feature_dim"]
+    step_count: Int[Array, ""]
 
 
 class ScaleDriftStream:

{alberta_framework-0.2.1.dist-info → alberta_framework-0.3.0.dist-info}/METADATA

@@ -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.dist-info → alberta_framework-0.3.0.dist-info}/RECORD

@@ -1,14 +1,14 @@
 alberta_framework/__init__.py,sha256=gAafDDmkivDdfnvDVff9zbVY9ilzqqfJ9KvpbRegKqs,5726
 alberta_framework/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 alberta_framework/core/__init__.py,sha256=PSrC4zSxgm_6YXWEQ80aZaunpbQ58QexxKmDDU-jp6c,522
-alberta_framework/core/learners.py,sha256=khZYkae5rlIyV13BW3-hrtPSjGFXPj2IUTM1z74xTTA,37724
-alberta_framework/core/normalizers.py,sha256=Z_d3H17qoXh87DE7k41imvWzkVJQ2xQgDUP7GYSNzAY,5903
-alberta_framework/core/optimizers.py,sha256=OefVuDDG1phh1QQIUyVPsQckl41VrpWFG7hY2eqyc64,14585
-alberta_framework/core/types.py,sha256=svV2Q5-0bj7reQ_hh-pRGp2wYfde5VWgTiRm4hUCDKI,9297
+alberta_framework/core/learners.py,sha256=gUhX7caXBfpWYgnvYTp5YKXfP6wbzB2T2gkSMMtrHDQ,38042
+alberta_framework/core/normalizers.py,sha256=QmKmha-mFgKi1KD-f8xuB2U175yQL6Ll0D4c8OONIl0,5927
+alberta_framework/core/optimizers.py,sha256=a4gYac5DyXReir9ycudRg8uQ9b53uLWTIldZ1A3Ae5c,14646
+alberta_framework/core/types.py,sha256=eqiPMVD8_QYJNg83rnE3XO9Z5BPczsg_LKh5dhmIgt4,9807
 alberta_framework/streams/__init__.py,sha256=bsDgWjWjotDQHMI2lno3dgk8N14pd-2mYAQpXAtCPx4,2035
 alberta_framework/streams/base.py,sha256=9rJxvUgmzd5u2bRV4vi5PxhUvj39EZTD4bZHo-Ptn-U,2168
 alberta_framework/streams/gymnasium.py,sha256=s733X7aEgy05hcSazjZEhBiJChtEL7uVpxwh0fXBQZA,21980
-alberta_framework/streams/synthetic.py,sha256=4R9GR7Kh0LT7GmGtPhzMJGr8HbhrAMUOjvPwLZk6nDg,32979
+alberta_framework/streams/synthetic.py,sha256=8njzQCFRi_iVgdPA3slyn46vFIHHkIwaZsABZyPwqnU,33507
 alberta_framework/utils/__init__.py,sha256=zfKfnbikhLp0J6UgVa8HeRo59gZHwqOc8jf03s7AaT4,2845
 alberta_framework/utils/experiments.py,sha256=ekGAzveCRgv9YZ5mfAD5Uf7h_PvQnxsNw2KeZN2eu00,10644
 alberta_framework/utils/export.py,sha256=W9RKfeTiyZcLColOGNjBfZU0N6QMXrfPn4pdYcm-OSk,15832
@@ -16,7 +16,7 @@ alberta_framework/utils/metrics.py,sha256=1cryNJoboO67vvRhausaucbYZFgdL_06vaf08U
 alberta_framework/utils/statistics.py,sha256=4fbzNlmsdUaM5lLW1BhL5B5MUpnqimQlwJklZ4x0y0U,15416
 alberta_framework/utils/timing.py,sha256=JOLq8CpCAV7LWOWkftxefduSFjaXnVwal1MFBKEMdJI,4049
 alberta_framework/utils/visualization.py,sha256=PmKBD3KGabNhgDizcNiGJEbVCyDL1YMUE5yTwgJHu2o,17924
-alberta_framework-0.2.1.dist-info/METADATA,sha256=pJ7SujFDZXrWxqjpSK8NFSBoYJpN5VfLneFnAmYG3hw,7763
-alberta_framework-0.2.1.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-alberta_framework-0.2.1.dist-info/licenses/LICENSE,sha256=TI1avodt5mvxz7sunyxIa0HlNgLQcmKNLeRjCVcgKmE,10754
-alberta_framework-0.2.1.dist-info/RECORD,,
+alberta_framework-0.3.0.dist-info/METADATA,sha256=OMn-jZdVSJ0vdLTLnpxzeTWx0UXlPFbGngA38wr_xDI,7872
+alberta_framework-0.3.0.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+alberta_framework-0.3.0.dist-info/licenses/LICENSE,sha256=TI1avodt5mvxz7sunyxIa0HlNgLQcmKNLeRjCVcgKmE,10754
+alberta_framework-0.3.0.dist-info/RECORD,,