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

alberta_framework/__init__.py

@@ -39,7 +39,7 @@ References
 - Tuning-free Step-size Adaptation (Mahmood et al., 2012)
 """
 
-__version__ = "0.1.0"
+__version__ = "0.2.0"
 
 # Core types
 # Learners
@@ -50,7 +50,9 @@ from alberta_framework.core.learners import (
     UpdateResult,
     metrics_to_dicts,
     run_learning_loop,
+    run_learning_loop_batched,
     run_normalized_learning_loop,
+    run_normalized_learning_loop_batched,
 )
 
 # Normalizers
@@ -64,6 +66,8 @@ from alberta_framework.core.normalizers import (
 from alberta_framework.core.optimizers import IDBD, LMS, Autostep, Optimizer
 from alberta_framework.core.types import (
     AutostepState,
+    BatchedLearningResult,
+    BatchedNormalizedResult,
     IDBDState,
     LearnerState,
     LMSState,
@@ -138,6 +142,8 @@ __all__ = [
     "__version__",
     # Types
     "AutostepState",
+    "BatchedLearningResult",
+    "BatchedNormalizedResult",
     "IDBDState",
     "LMSState",
     "LearnerState",
@@ -164,7 +170,9 @@ __all__ = [
     "NormalizedLearnerState",
     "NormalizedLinearLearner",
     "run_learning_loop",
+    "run_learning_loop_batched",
     "run_normalized_learning_loop",
+    "run_normalized_learning_loop_batched",
     "metrics_to_dicts",
     # Streams - protocol
     "ScanStream",

alberta_framework/core/learners.py

@@ -15,6 +15,8 @@ from alberta_framework.core.normalizers import NormalizerState, OnlineNormalizer
 from alberta_framework.core.optimizers import LMS, Optimizer
 from alberta_framework.core.types import (
     AutostepState,
+    BatchedLearningResult,
+    BatchedNormalizedResult,
     IDBDState,
     LearnerState,
     LMSState,
@@ -846,6 +848,196 @@ def run_normalized_learning_loop[StreamStateT](
         return final_learner, metrics
 
 
+def run_learning_loop_batched[StreamStateT](
+    learner: LinearLearner,
+    stream: ScanStream[StreamStateT],
+    num_steps: int,
+    keys: Array,
+    learner_state: LearnerState | None = None,
+    step_size_tracking: StepSizeTrackingConfig | None = None,
+) -> BatchedLearningResult:
+    """Run learning loop across multiple seeds in parallel using jax.vmap.
+
+    This function provides GPU parallelization for multi-seed experiments,
+    typically achieving 2-5x speedup over sequential execution.
+
+    Args:
+        learner: The learner to train
+        stream: Experience stream providing (observation, target) pairs
+        num_steps: Number of learning steps to run per seed
+        keys: JAX random keys with shape (num_seeds,) or (num_seeds, 2)
+        learner_state: Initial state (if None, will be initialized from stream).
+            The same initial state is used for all seeds.
+        step_size_tracking: Optional config for recording per-weight step-sizes.
+            When provided, history arrays have shape (num_seeds, num_recordings, ...)
+
+    Returns:
+        BatchedLearningResult containing:
+            - states: Batched final states with shape (num_seeds, ...) for each array
+            - metrics: Array of shape (num_seeds, num_steps, 3)
+            - step_size_history: Batched history or None if tracking disabled
+
+    Examples:
+    ```python
+    import jax.random as jr
+    from alberta_framework import LinearLearner, IDBD, RandomWalkStream
+    from alberta_framework import run_learning_loop_batched
+
+    stream = RandomWalkStream(feature_dim=10)
+    learner = LinearLearner(optimizer=IDBD())
+
+    # Run 30 seeds in parallel
+    keys = jr.split(jr.key(42), 30)
+    result = run_learning_loop_batched(learner, stream, num_steps=10000, keys=keys)
+
+    # result.metrics has shape (30, 10000, 3)
+    mean_error = result.metrics[:, :, 0].mean(axis=0)  # Average over seeds
+    ```
+    """
+    # Define single-seed function that returns consistent structure
+    def single_seed_run(key: Array) -> tuple[LearnerState, Array, StepSizeHistory | None]:
+        result = run_learning_loop(
+            learner, stream, num_steps, key, learner_state, step_size_tracking
+        )
+        if step_size_tracking is not None:
+            state, metrics, history = result
+            return state, metrics, history
+        else:
+            state, metrics = result
+            # Return None for history to maintain consistent output structure
+            return state, metrics, None
+
+    # vmap over the keys dimension
+    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:
+        batched_step_size_history = StepSizeHistory(
+            step_sizes=batched_history.step_sizes,
+            bias_step_sizes=batched_history.bias_step_sizes,
+            recording_indices=batched_history.recording_indices,
+            normalizers=batched_history.normalizers,
+        )
+    else:
+        batched_step_size_history = None
+
+    return BatchedLearningResult(
+        states=batched_states,
+        metrics=batched_metrics,
+        step_size_history=batched_step_size_history,
+    )
+
+
+def run_normalized_learning_loop_batched[StreamStateT](
+    learner: NormalizedLinearLearner,
+    stream: ScanStream[StreamStateT],
+    num_steps: int,
+    keys: Array,
+    learner_state: NormalizedLearnerState | None = None,
+    step_size_tracking: StepSizeTrackingConfig | None = None,
+    normalizer_tracking: NormalizerTrackingConfig | None = None,
+) -> BatchedNormalizedResult:
+    """Run normalized learning loop across multiple seeds in parallel using jax.vmap.
+
+    This function provides GPU parallelization for multi-seed experiments with
+    normalized learners, typically achieving 2-5x speedup over sequential execution.
+
+    Args:
+        learner: The normalized learner to train
+        stream: Experience stream providing (observation, target) pairs
+        num_steps: Number of learning steps to run per seed
+        keys: JAX random keys with shape (num_seeds,) or (num_seeds, 2)
+        learner_state: Initial state (if None, will be initialized from stream).
+            The same initial state is used for all seeds.
+        step_size_tracking: Optional config for recording per-weight step-sizes.
+            When provided, history arrays have shape (num_seeds, num_recordings, ...)
+        normalizer_tracking: Optional config for recording normalizer state.
+            When provided, history arrays have shape (num_seeds, num_recordings, ...)
+
+    Returns:
+        BatchedNormalizedResult containing:
+            - states: Batched final states with shape (num_seeds, ...) for each array
+            - metrics: Array of shape (num_seeds, num_steps, 4)
+            - step_size_history: Batched history or None if tracking disabled
+            - normalizer_history: Batched history or None if tracking disabled
+
+    Examples:
+    ```python
+    import jax.random as jr
+    from alberta_framework import NormalizedLinearLearner, IDBD, RandomWalkStream
+    from alberta_framework import run_normalized_learning_loop_batched
+
+    stream = RandomWalkStream(feature_dim=10)
+    learner = NormalizedLinearLearner(optimizer=IDBD())
+
+    # Run 30 seeds in parallel
+    keys = jr.split(jr.key(42), 30)
+    result = run_normalized_learning_loop_batched(
+        learner, stream, num_steps=10000, keys=keys
+    )
+
+    # result.metrics has shape (30, 10000, 4)
+    mean_error = result.metrics[:, :, 0].mean(axis=0)  # Average over seeds
+    ```
+    """
+    # Define single-seed function that returns consistent structure
+    def single_seed_run(
+        key: Array,
+    ) -> tuple[
+        NormalizedLearnerState, Array, StepSizeHistory | None, NormalizerHistory | None
+    ]:
+        result = run_normalized_learning_loop(
+            learner, stream, num_steps, key, learner_state,
+            step_size_tracking, normalizer_tracking
+        )
+
+        # 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
+            return state, metrics, ss_history, norm_history
+        elif step_size_tracking is not None:
+            state, metrics, ss_history = result
+            return state, metrics, ss_history, None
+        elif normalizer_tracking is not None:
+            state, metrics, norm_history = result
+            return state, metrics, None, norm_history
+        else:
+            state, metrics = result
+            return state, metrics, None, None
+
+    # vmap over the keys dimension
+    batched_states, batched_metrics, batched_ss_history, batched_norm_history = (
+        jax.vmap(single_seed_run)(keys)
+    )
+
+    # Reconstruct batched histories if tracking was enabled
+    if step_size_tracking is not None:
+        batched_step_size_history = StepSizeHistory(
+            step_sizes=batched_ss_history.step_sizes,
+            bias_step_sizes=batched_ss_history.bias_step_sizes,
+            recording_indices=batched_ss_history.recording_indices,
+            normalizers=batched_ss_history.normalizers,
+        )
+    else:
+        batched_step_size_history = None
+
+    if normalizer_tracking is not None:
+        batched_normalizer_history = NormalizerHistory(
+            means=batched_norm_history.means,
+            variances=batched_norm_history.variances,
+            recording_indices=batched_norm_history.recording_indices,
+        )
+    else:
+        batched_normalizer_history = None
+
+    return BatchedNormalizedResult(
+        states=batched_states,
+        metrics=batched_metrics,
+        step_size_history=batched_step_size_history,
+        normalizer_history=batched_normalizer_history,
+    )
+
+
 def metrics_to_dicts(metrics: Array, normalized: bool = False) -> list[dict[str, float]]:
     """Convert metrics array to list of dicts for backward compatibility.
 

alberta_framework/core/types.py

@@ -4,11 +4,14 @@ This module defines the core data types used throughout the framework,
 following JAX conventions with immutable NamedTuples for state management.
 """
 
-from typing import NamedTuple
+from typing import TYPE_CHECKING, NamedTuple
 
 import jax.numpy as jnp
 from jax import Array
 
+if TYPE_CHECKING:
+    from alberta_framework.core.learners import NormalizedLearnerState
+
 # Type aliases for clarity
 Observation = Array  # x_t: feature vector
 Target = Array  # y*_t: desired output
@@ -164,6 +167,45 @@ class NormalizerHistory(NamedTuple):
     recording_indices: Array  # (num_recordings,)
 
 
+class BatchedLearningResult(NamedTuple):
+    """Result from batched learning loop across multiple seeds.
+
+    Used with `run_learning_loop_batched` for vmap-based GPU parallelization.
+
+    Attributes:
+        states: Batched learner states - each array has shape (num_seeds, ...)
+        metrics: Metrics array with shape (num_seeds, num_steps, 3)
+            where columns are [squared_error, error, mean_step_size]
+        step_size_history: Optional step-size history with batched shapes,
+            or None if tracking was disabled
+    """
+
+    states: "LearnerState"  # Batched: each array has shape (num_seeds, ...)
+    metrics: Array  # Shape: (num_seeds, num_steps, 3)
+    step_size_history: StepSizeHistory | None
+
+
+class BatchedNormalizedResult(NamedTuple):
+    """Result from batched normalized learning loop across multiple seeds.
+
+    Used with `run_normalized_learning_loop_batched` for vmap-based GPU parallelization.
+
+    Attributes:
+        states: Batched normalized learner states - each array has shape (num_seeds, ...)
+        metrics: Metrics array with shape (num_seeds, num_steps, 4)
+            where columns are [squared_error, error, mean_step_size, normalizer_mean_var]
+        step_size_history: Optional step-size history with batched shapes,
+            or None if tracking was disabled
+        normalizer_history: Optional normalizer history with batched shapes,
+            or None if tracking was disabled
+    """
+
+    states: "NormalizedLearnerState"  # Batched: each array has shape (num_seeds, ...)
+    metrics: Array  # Shape: (num_seeds, num_steps, 4)
+    step_size_history: StepSizeHistory | None
+    normalizer_history: NormalizerHistory | None
+
+
 def create_lms_state(step_size: float = 0.01) -> LMSState:
     """Create initial LMS optimizer state.
 

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: alberta-framework
-Version: 0.1.1
+Version: 0.2.1
 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

{alberta_framework-0.1.1.dist-info → alberta_framework-0.2.1.dist-info}/RECORD

@@ -1,10 +1,10 @@
-alberta_framework/__init__.py,sha256=LUrsm6WFh5-Mxg78d1G-Qe015nkGgcCDhSw5lf3UkFo,5460
+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=dnRQ5B16oGYpamDJIRYzR54ED9bvW0lpa8c_suC6YBA,29879
+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=mtpVEr2qJ0XzZyjOsUdChmS7T7mrXBDMHb-jfkrT9JY,7503
+alberta_framework/core/types.py,sha256=svV2Q5-0bj7reQ_hh-pRGp2wYfde5VWgTiRm4hUCDKI,9297
 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
@@ -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.1.1.dist-info/METADATA,sha256=Ny-LxHiqZVNXZbu5f8ZyBSLCEZd2KsBhA9iROV7tNiU,7763
-alberta_framework-0.1.1.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-alberta_framework-0.1.1.dist-info/licenses/LICENSE,sha256=TI1avodt5mvxz7sunyxIa0HlNgLQcmKNLeRjCVcgKmE,10754
-alberta_framework-0.1.1.dist-info/RECORD,,
+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,,