marn 0.1.0__tar.gz

marn-0.1.0/LICENSE

@@ -0,0 +1,161 @@
+Apache License
+Version 2.0, January 2004
+http://www.apache.org/licenses/
+
+TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+1. Definitions.
+
+"License" shall mean the terms and conditions for use, reproduction, and
+distribution as defined by Sections 1 through 9 of this document.
+
+"Licensor" shall mean the copyright owner or entity authorized by the
+copyright owner that is granting the License.
+
+"Legal Entity" shall mean the union of the acting entity and all other
+entities that control, are controlled by, or are under common control with
+that entity. For the purposes of this definition, "control" means (i) the
+power, direct or indirect, to cause the direction or management of such
+entity, whether by contract or otherwise, or (ii) ownership of fifty percent
+(50%) or more of the outstanding shares, or (iii) beneficial ownership of
+such entity.
+
+"You" (or "Your") shall mean an individual or Legal Entity exercising
+permissions granted by this License.
+
+"Source" form shall mean the preferred form for making modifications,
+including but not limited to software source code, documentation source, and
+configuration files.
+
+"Object" form shall mean any form resulting from mechanical transformation
+or translation of a Source form, including but not limited to compiled object
+code, generated documentation, and conversions to other media types.
+
+"Work" shall mean the work of authorship, whether in Source or Object form,
+made available under the License, as indicated by a copyright notice that is
+included in or attached to the work.
+
+"Derivative Works" shall mean any work, whether in Source or Object form,
+that is based on (or derived from) the Work and for which the editorial
+revisions, annotations, elaborations, or other modifications represent, as a
+whole, an original work of authorship. For the purposes of this License,
+Derivative Works shall not include works that remain separable from, or
+merely link (or bind by name) to the interfaces of, the Work and Derivative
+Works thereof.
+
+"Contribution" shall mean any work of authorship, including the original
+version of the Work and any modifications or additions to that Work or
+Derivative Works thereof, that is intentionally submitted to Licensor for
+inclusion in the Work by the copyright owner or by an individual or Legal
+Entity authorized to submit on behalf of the copyright owner. For the purposes
+of this definition, "submitted" means any form of electronic, verbal, or
+written communication sent to the Licensor or its representatives, including
+but not limited to communication on electronic mailing lists, source code
+control systems, and issue tracking systems that are managed by, or on behalf
+of, the Licensor for the purpose of discussing and improving the Work, but
+excluding communication that is conspicuously marked or otherwise designated
+in writing by the copyright owner as "Not a Contribution."
+
+"Contributor" shall mean Licensor and any individual or Legal Entity on behalf
+of whom a Contribution has been received by Licensor and subsequently
+incorporated within the Work.
+
+2. Grant of Copyright License. Subject to the terms and conditions of this
+License, each Contributor hereby grants to You a perpetual, worldwide,
+non-exclusive, no-charge, royalty-free, irrevocable copyright license to
+reproduce, prepare Derivative Works of, publicly display, publicly perform,
+sublicense, and distribute the Work and such Derivative Works in Source or
+Object form.
+
+3. Grant of Patent License. Subject to the terms and conditions of this
+License, each Contributor hereby grants to You a perpetual, worldwide,
+non-exclusive, no-charge, royalty-free, irrevocable (except as stated in this
+section) patent license to make, have made, use, offer to sell, sell, import,
+and otherwise transfer the Work, where such license applies only to those
+patent claims licensable by such Contributor that are necessarily infringed by
+their Contribution(s) alone or by combination of their Contribution(s) with
+the Work to which such Contribution(s) was submitted. If You institute patent
+litigation against any entity (including a cross-claim or counterclaim in a
+lawsuit) alleging that the Work or a Contribution incorporated within the Work
+constitutes direct or contributory patent infringement, then any patent
+licenses granted to You under this License for that Work shall terminate as of
+the date such litigation is filed.
+
+4. Redistribution. You may reproduce and distribute copies of the Work or
+Derivative Works thereof in any medium, with or without modifications, and in
+Source or Object form, provided that You meet the following conditions:
+
+(a) You must give any other recipients of the Work or Derivative Works a copy
+of this License; and
+
+(b) You must cause any modified files to carry prominent notices stating that
+You changed the files; and
+
+(c) You must retain, in the Source form of any Derivative Works that You
+distribute, all copyright, patent, trademark, and attribution notices from the
+Source form of the Work, excluding those notices that do not pertain to any
+part of the Derivative Works; and
+
+(d) If the Work includes a "NOTICE" text file as part of its distribution,
+then any Derivative Works that You distribute must include a readable copy of
+the attribution notices contained within such NOTICE file, excluding those
+notices that do not pertain to any part of the Derivative Works, in at least
+one of the following places: within a NOTICE text file distributed as part of
+the Derivative Works; within the Source form or documentation, if provided
+along with the Derivative Works; or within a display generated by the
+Derivative Works, if and wherever such third-party notices normally appear.
+The contents of the NOTICE file are for informational purposes only and do not
+modify the License. You may add Your own attribution notices within Derivative
+Works that You distribute, alongside or as an addendum to the NOTICE text from
+the Work, provided that such additional attribution notices cannot be construed
+as modifying the License.
+
+You may add Your own copyright statement to Your modifications and may provide
+additional or different license terms and conditions for use, reproduction, or
+distribution of Your modifications, or for any such Derivative Works as a
+whole, provided Your use, reproduction, and distribution of the Work otherwise
+complies with the conditions stated in this License.
+
+5. Submission of Contributions. Unless You explicitly state otherwise, any
+Contribution intentionally submitted for inclusion in the Work by You to the
+Licensor shall be under the terms and conditions of this License, without any
+additional terms or conditions. Notwithstanding the above, nothing herein
+shall supersede or modify the terms of any separate license agreement you may
+have executed with Licensor regarding such Contributions.
+
+6. Trademarks. This License does not grant permission to use the trade names,
+trademarks, service marks, or product names of the Licensor, except as
+required for reasonable and customary use in describing the origin of the Work
+and reproducing the content of the NOTICE file.
+
+7. Disclaimer of Warranty. Unless required by applicable law or agreed to in
+writing, Licensor provides the Work (and each Contributor provides its
+Contributions) on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+KIND, either express or implied, including, without limitation, any warranties
+or conditions of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+PARTICULAR PURPOSE. You are solely responsible for determining the
+appropriateness of using or redistributing the Work and assume any risks
+associated with Your exercise of permissions under this License.
+
+8. Limitation of Liability. In no event and under no legal theory, whether in
+tort (including negligence), contract, or otherwise, unless required by
+applicable law (such as deliberate and grossly negligent acts) or agreed to in
+writing, shall any Contributor be liable to You for damages, including any
+direct, indirect, special, incidental, or consequential damages of any
+character arising as a result of this License or out of the use or inability
+to use the Work (including but not limited to damages for loss of goodwill,
+work stoppage, computer failure or malfunction, or any and all other
+commercial damages or losses), even if such Contributor has been advised of
+the possibility of such damages.
+
+9. Accepting Warranty or Additional Liability. While redistributing the Work
+or Derivative Works thereof, You may choose to offer, and charge a fee for,
+acceptance of support, warranty, indemnity, or other liability obligations
+and/or rights consistent with this License. However, in accepting such
+obligations, You may act only on Your own behalf and on Your sole
+responsibility, not on behalf of any other Contributor, and only if You agree
+to indemnify, defend, and hold each Contributor harmless for any liability
+incurred by, or claims asserted against, such Contributor by reason of your
+accepting any such warranty or additional liability.
+
+END OF TERMS AND CONDITIONS

marn-0.1.0/PKG-INFO

@@ -0,0 +1,125 @@
+Metadata-Version: 2.4
+Name: marn
+Version: 0.1.0
+Summary: 
+License: Apache-2.0
+License-File: LICENSE
+Author: Arjun Manjunath
+Author-email: dev.arjunmnath@gmail.com
+Requires-Python: >=3.12, !=2.7.*, !=3.0.*, !=3.1.*, !=3.2.*, !=3.3.*, !=3.4.*, !=3.5.*, !=3.6.*, !=3.7.*, !=3.8.*, !=3.9.*, !=3.10.*, !=3.11.*
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: 3.14
+Requires-Dist: accelerate (>=0.30)
+Requires-Dist: matplotlib (>=3.11.2,<4.0.0)
+Requires-Dist: pydantic (>=2.0)
+Requires-Dist: pyyaml (>=6.0.3,<7.0.0)
+Requires-Dist: scikit-learn (>=1.9.0,<2.0.0)
+Requires-Dist: torch (>=2.6,<3.0)
+Requires-Dist: torchvision (>=0.27.1,<0.28.0)
+Description-Content-Type: text/markdown
+
+# Manifold Regularized Networks (MaRN)
+
+[![License](https://img.shields.io/badge/License-Apache%202.0-blue.svg)](https://opensource.org/licenses/Apache-2.0)
+[![Python Version](https://img.shields.io/badge/python-3.12-blue)](https://www.python.org/downloads/)
+[![PyTorch Version](https://img.shields.io/badge/PyTorch-%3E%3D2.6%20%7C%20%3C3.0-orange)](https://pytorch.org/)
+[![Build Status](https://img.shields.io/badge/build-passing-brightgreen)](#)
+[![Documentation Status](https://readthedocs.org/projects/mapping-networks/badge/?version=latest)](https://mapping-networks.readthedocs.io/en/latest/?badge=latest)
+[![arXiv](https://img.shields.io/badge/arXiv-2602.19134-b31b1b.svg)](https://arxiv.org/abs/2602.19134)
+
+`marn` is a model-agnostic, production-oriented PyTorch package for training target models through low-dimensional parameter manifolds. It is inspired by the paper [**Mapping Networks**](https://doi.org/10.48550/arXiv.2602.19134). The library decouples model architecture from parameter representation, allowing you to optimize neural networks by updating compact, trainable latent coordinates instead of mutating the target module's parameters.
+
+---
+
+## Installation
+
+Install `marn` from your local workspace:
+
+```bash
+pip install marn
+```
+
+Or add it using Poetry:
+
+```bash
+poetry add marn
+```
+
+---
+
+## Quickstart
+
+Train a standard PyTorch model through a low-dimensional layer-wise latent representation:
+
+```python
+import torch
+from torch import nn
+from torch.utils.data import DataLoader, TensorDataset
+from marn import MappingModel, MappingLoss, ClassificationLoss, MappingTrainer
+
+# 1. Create target model and training data
+target = nn.Sequential(nn.Linear(10, 16), nn.ReLU(), nn.Linear(16, 2))
+data = TensorDataset(torch.randn(100, 10), torch.randint(0, 2, (100,)))
+loader = DataLoader(data, batch_size=16)
+
+# 2. Wrap target model with low-dimensional mapping strategy
+model = MappingModel(
+    target_model=target,
+    latent_dim=32,
+    strategy="layerwise"
+)
+
+# 3. Configure composite loss and trainer
+loss_fn = MappingLoss(task_loss=ClassificationLoss())
+trainer = MappingTrainer(
+    model=model,
+    train_loader=loader,
+    loss_fn=loss_fn,
+    learning_rate=1e-3,
+)
+
+# 4. Train the latent parameters
+trainer.fit(epochs=5)
+```
+
+---
+
+## Core Concepts
+
+- **Latent Manifold**: Trainable coordinates $z$ which undergo transformation to map to target parameter space.
+- **BaseMapper**: Projects $z$ to flat generated parameter descriptors (e.g., using fixed, orthogonal projections).
+- **BaseModulation**: Integrates generated descriptors back into target parameters (e.g., additive $W_{ij} \leftarrow W_{ij} + \alpha z_i$, or affine).
+- **Generation Strategy**: Defines the mapping scope. 
+  - `"slvt"` (Single Latent Vector Training) projects the entire model from a single global latent vector.
+  - `"layerwise"` constructs independent smaller latent coordinates per layer.
+  - `"grouped"` allows custom parameter subdivision.
+- **MappingLoss**: A composite loss that balances task loss (classification/regression) with stability, smoothness, and cosine alignment regularization components.
+
+---
+
+## Paper Correspondence
+
+The code maps directly to the concepts defined in the paper:
+- **Fixed Projection Matrices**: Registered as PyTorch buffers inside `MLPMapper` so they stay frozen and are excluded from DDP/optimizer updates.
+- **Additive Modulation**: Implemented in `AdditiveModulation` representing $w_{ij} \leftarrow w_{ij} + \alpha \cdot z_i$.
+- **Regularization Terms**: Fully implemented in `MappingLoss`:
+  - **Stability Loss** ($L_{\text{stability}}$): penalizes changes in output predictions when adding small noise perturbation to the latent vector (`StabilityLoss`).
+  - **Smoothness Loss** ($L_{\text{smoothness}}$): penalizes the Jacobian norm of the mapper to enforce a smooth manifold (`SmoothnessLoss`).
+  - **Alignment Loss** ($L_{\text{alignment}}$): maximizes alignment via cosine distance between latent vectors and weight summaries (`AlignmentLoss`).
+
+---
+
+## Documentation & cookbook
+
+- **Docs**: [mapping-networks.readthedocs.io](https://mapping-networks.readthedocs.io/en/latest/)
+- **User guide & API**: organized under the PyData Sphinx theme in `docs/`
+- **Cookbook**: eight runnable scripts in [`cookbook/`](cookbook/) (also summarized in the [Cookbook docs](https://mapping-networks.readthedocs.io/en/latest/cookbook/index.html))
+
+```bash
+poetry run python cookbook/01_single_latent_classification.py
+```
+
+

marn-0.1.0/README.md

@@ -0,0 +1,101 @@
+# Manifold Regularized Networks (MaRN)
+
+[![License](https://img.shields.io/badge/License-Apache%202.0-blue.svg)](https://opensource.org/licenses/Apache-2.0)
+[![Python Version](https://img.shields.io/badge/python-3.12-blue)](https://www.python.org/downloads/)
+[![PyTorch Version](https://img.shields.io/badge/PyTorch-%3E%3D2.6%20%7C%20%3C3.0-orange)](https://pytorch.org/)
+[![Build Status](https://img.shields.io/badge/build-passing-brightgreen)](#)
+[![Documentation Status](https://readthedocs.org/projects/mapping-networks/badge/?version=latest)](https://mapping-networks.readthedocs.io/en/latest/?badge=latest)
+[![arXiv](https://img.shields.io/badge/arXiv-2602.19134-b31b1b.svg)](https://arxiv.org/abs/2602.19134)
+
+`marn` is a model-agnostic, production-oriented PyTorch package for training target models through low-dimensional parameter manifolds. It is inspired by the paper [**Mapping Networks**](https://doi.org/10.48550/arXiv.2602.19134). The library decouples model architecture from parameter representation, allowing you to optimize neural networks by updating compact, trainable latent coordinates instead of mutating the target module's parameters.
+
+---
+
+## Installation
+
+Install `marn` from your local workspace:
+
+```bash
+pip install marn
+```
+
+Or add it using Poetry:
+
+```bash
+poetry add marn
+```
+
+---
+
+## Quickstart
+
+Train a standard PyTorch model through a low-dimensional layer-wise latent representation:
+
+```python
+import torch
+from torch import nn
+from torch.utils.data import DataLoader, TensorDataset
+from marn import MappingModel, MappingLoss, ClassificationLoss, MappingTrainer
+
+# 1. Create target model and training data
+target = nn.Sequential(nn.Linear(10, 16), nn.ReLU(), nn.Linear(16, 2))
+data = TensorDataset(torch.randn(100, 10), torch.randint(0, 2, (100,)))
+loader = DataLoader(data, batch_size=16)
+
+# 2. Wrap target model with low-dimensional mapping strategy
+model = MappingModel(
+    target_model=target,
+    latent_dim=32,
+    strategy="layerwise"
+)
+
+# 3. Configure composite loss and trainer
+loss_fn = MappingLoss(task_loss=ClassificationLoss())
+trainer = MappingTrainer(
+    model=model,
+    train_loader=loader,
+    loss_fn=loss_fn,
+    learning_rate=1e-3,
+)
+
+# 4. Train the latent parameters
+trainer.fit(epochs=5)
+```
+
+---
+
+## Core Concepts
+
+- **Latent Manifold**: Trainable coordinates $z$ which undergo transformation to map to target parameter space.
+- **BaseMapper**: Projects $z$ to flat generated parameter descriptors (e.g., using fixed, orthogonal projections).
+- **BaseModulation**: Integrates generated descriptors back into target parameters (e.g., additive $W_{ij} \leftarrow W_{ij} + \alpha z_i$, or affine).
+- **Generation Strategy**: Defines the mapping scope. 
+  - `"slvt"` (Single Latent Vector Training) projects the entire model from a single global latent vector.
+  - `"layerwise"` constructs independent smaller latent coordinates per layer.
+  - `"grouped"` allows custom parameter subdivision.
+- **MappingLoss**: A composite loss that balances task loss (classification/regression) with stability, smoothness, and cosine alignment regularization components.
+
+---
+
+## Paper Correspondence
+
+The code maps directly to the concepts defined in the paper:
+- **Fixed Projection Matrices**: Registered as PyTorch buffers inside `MLPMapper` so they stay frozen and are excluded from DDP/optimizer updates.
+- **Additive Modulation**: Implemented in `AdditiveModulation` representing $w_{ij} \leftarrow w_{ij} + \alpha \cdot z_i$.
+- **Regularization Terms**: Fully implemented in `MappingLoss`:
+  - **Stability Loss** ($L_{\text{stability}}$): penalizes changes in output predictions when adding small noise perturbation to the latent vector (`StabilityLoss`).
+  - **Smoothness Loss** ($L_{\text{smoothness}}$): penalizes the Jacobian norm of the mapper to enforce a smooth manifold (`SmoothnessLoss`).
+  - **Alignment Loss** ($L_{\text{alignment}}$): maximizes alignment via cosine distance between latent vectors and weight summaries (`AlignmentLoss`).
+
+---
+
+## Documentation & cookbook
+
+- **Docs**: [mapping-networks.readthedocs.io](https://mapping-networks.readthedocs.io/en/latest/)
+- **User guide & API**: organized under the PyData Sphinx theme in `docs/`
+- **Cookbook**: eight runnable scripts in [`cookbook/`](cookbook/) (also summarized in the [Cookbook docs](https://mapping-networks.readthedocs.io/en/latest/cookbook/index.html))
+
+```bash
+poetry run python cookbook/01_single_latent_classification.py
+```
+

marn-0.1.0/pyproject.toml

@@ -0,0 +1,58 @@
+[tool.poetry]
+name = "marn"
+version = "0.1.0"
+description = ""
+authors = ["Arjun Manjunath <dev.arjunmnath@gmail.com>"]
+readme = "README.md"
+license = "Apache-2.0"
+
+packages = [
+    { include = "marn", from = "src" }
+]
+
+[tool.poetry.dependencies]
+python = ">=3.12,<3.14.1 || >3.14.1,<3.15"
+torch = ">=2.6,<3.0"
+pydantic = ">=2.0"
+pyyaml = "^6.0.3"
+torchvision = "^0.27.1"
+accelerate = ">=0.30"
+scikit-learn = "^1.9.0"
+matplotlib = "^3.11.2"
+
+[tool.poetry.group.dev.dependencies]
+pytest = "^8.3"
+pytest-cov = "^6.0"
+ruff = "^0.12"
+mypy = "^1.16"
+pre-commit = "^4.2"
+types-pyyaml = "^6.0.12.20260518"
+
+
+[tool.poetry.group.docs.dependencies]
+sphinx = "^9.1.0"
+myst-parser = "^5.1.0"
+pydata-sphinx-theme = "^0.15.2"
+sphinx-autobuild = "^2025.8.25"
+
+
+[tool.ruff]
+line-length = 100
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+
+[tool.mypy]
+strict = true
+
+[[tool.mypy.overrides]]
+module = [
+    "matplotlib.*",
+    "torchvision.*",
+    "accelerate.*",
+]
+ignore_missing_imports = true
+
+[build-system]
+requires = ["poetry-core"]
+build-backend = "poetry.core.masonry.api"

marn-0.1.0/src/marn/__init__.py

@@ -0,0 +1,154 @@
+"""Train PyTorch models through low-dimensional parameter mappings."""
+
+from marn.generators import (
+    GroupedGenerator,
+    LayerGenerator,
+    LazyLayerwiseGenerator,
+    LayerwiseGenerator,
+    ParameterGenerator,
+    SingleVectorGenerator,
+    FineTuningGenerator,
+    LRDGenerator,
+)
+from marn.distributed import (
+    benchmark_strategies,
+    cleanup_ddp,
+    is_ddp_available,
+    profile_peak_memory,
+    setup_ddp,
+    wrap_ddp,
+)
+from marn.models import (
+    ForwardResult,
+    MappingModel,
+    TargetModel,
+    UnsupportedTargetModelError,
+)
+from marn.losses import (
+    AlignmentLoss,
+    BaseLoss,
+    ClassificationLoss,
+    LossOutput,
+    MappingLoss,
+    RegressionLoss,
+    SmoothnessLoss,
+    StabilityLoss,
+    TaskLoss,
+    TrainingContext,
+)
+from marn.mappers import BaseMapper, MLPMapper, ResidualMLPMapper
+from marn.modulation import (
+    AdditiveModulation,
+    AffineModulation,
+    BaseModulation,
+    LowRankModulation,
+)
+from marn.runtime import ParameterEntry, ParameterSpec, ParameterTree
+from marn.strategies import (
+    GroupedStrategy,
+    LayerwiseStrategy,
+    SLVTStrategy,
+    FineTuningStrategy,
+    LRDStrategy,
+)
+from marn.config import (
+    GeneratorConfig,
+    LossConfig,
+    MapperConfig,
+    MappingConfig,
+    TaskLossConfig,
+    TrainerConfig,
+    load_config,
+)
+from marn.registry import (
+    GENERATOR_REGISTRY,
+    LOSS_REGISTRY,
+    MAPPER_REGISTRY,
+    MODULATION_REGISTRY,
+    Registry,
+)
+from marn.callbacks import Callback, EarlyStopping, MetricLogger
+from marn.trainers import (
+    BatchAdapter,
+    MappingBatchAdapter,
+    MappingTrainer,
+    TupleBatchAdapter,
+    LRFinderResult,
+)
+from marn.checkpoint import (
+    CheckpointCompatibilityError,
+    CheckpointSchema,
+    load_checkpoint,
+    save_checkpoint,
+)
+
+__all__ = [
+    "AdditiveModulation",
+    "AffineModulation",
+    "AlignmentLoss",
+    "BaseLoss",
+    "BaseMapper",
+    "BaseModulation",
+    "BatchAdapter",
+    "Callback",
+    "CheckpointCompatibilityError",
+    "CheckpointSchema",
+    "ClassificationLoss",
+    "EarlyStopping",
+    "ForwardResult",
+    "FineTuningGenerator",
+    "GENERATOR_REGISTRY",
+    "GeneratorConfig",
+    "GroupedGenerator",
+    "GroupedStrategy",
+    "LOSS_REGISTRY",
+    "LayerGenerator",
+    "LazyLayerwiseGenerator",
+    "LayerwiseGenerator",
+    "LayerwiseStrategy",
+    "LRDGenerator",
+    "LossConfig",
+    "LossOutput",
+    "LowRankModulation",
+    "MAPPER_REGISTRY",
+    "MODULATION_REGISTRY",
+    "MapperConfig",
+    "MappingBatchAdapter",
+    "MappingConfig",
+    "MappingLoss",
+    "MappingModel",
+    "MappingTrainer",
+    "MLPMapper",
+    "MetricLogger",
+    "ParameterEntry",
+    "ParameterGenerator",
+    "ParameterSpec",
+    "ParameterTree",
+    "RegressionLoss",
+    "Registry",
+    "ResidualMLPMapper",
+    "SLVTStrategy",
+    "SingleVectorGenerator",
+    "FineTuningStrategy",
+    "LRDStrategy",
+    "LRFinderResult",
+    "SmoothnessLoss",
+    "StabilityLoss",
+    "TargetModel",
+    "TaskLoss",
+    "TaskLossConfig",
+    "TrainerConfig",
+    "TrainingContext",
+    "TupleBatchAdapter",
+    "UnsupportedTargetModelError",
+    "benchmark_strategies",
+    "cleanup_ddp",
+    "is_ddp_available",
+    "load_checkpoint",
+    "load_config",
+    "profile_peak_memory",
+    "save_checkpoint",
+    "setup_ddp",
+    "wrap_ddp",
+]
+__version__ = "0.1.0"

marn-0.1.0/src/marn/callbacks/__init__.py

@@ -0,0 +1,7 @@
+"""Training lifecycle callbacks."""
+
+from marn.callbacks.base import Callback
+from marn.callbacks.early_stopping import EarlyStopping
+from marn.callbacks.logger import MetricLogger
+
+__all__ = ["Callback", "EarlyStopping", "MetricLogger"]

marn-0.1.0/src/marn/callbacks/base.py

@@ -0,0 +1,51 @@
+"""Observer callback base class for training lifecycle hooks."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from marn.losses.outputs import LossOutput
+
+
+class Callback:
+    """Base class for training lifecycle callbacks.
+
+    All hooks are no-ops by default — subclass and override the ones you need.
+    Callbacks receive the trainer instance so they can inspect or modify state
+    (e.g. set ``trainer.should_stop = True`` for early termination).
+
+    Hook execution order within each event:
+
+    1. ``on_fit_start``
+    2. For each epoch:
+       a. ``on_epoch_start``
+       b. For each batch: ``on_batch_start`` → ``on_batch_end``
+       c. ``on_validation_start`` → ``on_validation_end``
+       d. ``on_epoch_end``
+    3. ``on_fit_end``
+    """
+
+    def on_fit_start(self, trainer: Any) -> None:
+        """Called once at the beginning of :meth:`fit`."""
+
+    def on_fit_end(self, trainer: Any) -> None:
+        """Called once at the end of :meth:`fit`."""
+
+    def on_epoch_start(self, trainer: Any, epoch: int) -> None:
+        """Called at the beginning of each epoch."""
+
+    def on_epoch_end(self, trainer: Any, epoch: int, metrics: dict[str, float]) -> None:
+        """Called at the end of each epoch with aggregated metrics."""
+
+    def on_batch_start(self, trainer: Any, batch_idx: int) -> None:
+        """Called before processing each batch."""
+
+    def on_batch_end(self, trainer: Any, batch_idx: int, loss_output: LossOutput) -> None:
+        """Called after processing each batch with the loss output."""
+
+    def on_validation_start(self, trainer: Any) -> None:
+        """Called before the validation loop."""
+
+    def on_validation_end(self, trainer: Any, metrics: dict[str, float]) -> None:
+        """Called after the validation loop with validation metrics."""