tsn-affinity

API reference

The API reference is split into per-package pages so each section stays focused. Use the navigation bar to jump to a specific package, or browse from here.

Public surface

TSN-Affinity exposes the following public packages. Every symbol listed below is also re-exported from the top-level tsn_affinity namespace, so from tsn_affinity import DecisionTransformer works without navigating the package layout.

- **`tsn_affinity.core`** — Model definitions, configs, exceptions, logging. See [Core](/tsn-affinity/api/core.html). - **`tsn_affinity.sparse`** — Sparse layers and conversion utilities. See [Sparse layers](/tsn-affinity/api/sparse.html). - **`tsn_affinity.routing`** — Affinity metrics and warm-starting. See [Routing](/tsn-affinity/api/routing.html). - **`tsn_affinity.strategies`** — Continual learning strategies. See [Strategies](/tsn-affinity/api/strategies.html). - **`tsn_affinity.data`** — Trajectory schemas, loaders, batch generation. See [Data](/tsn-affinity/api/data.html). - **`tsn_affinity.benchmarks`** — Adapters, registries, metrics, Atari baselines. See [Benchmarks](/tsn-affinity/api/benchmarks.html). - **`tsn_affinity.interfaces`** — Abstract protocols. See [Interfaces](/tsn-affinity/api/interfaces.html). - **`tsn_affinity.run`** — Result analysis. Covered in [Benchmarks](/tsn-affinity/api/benchmarks.html). - **`tsn_affinity.services`** — High-level orchestration. Covered in [Strategies](/tsn-affinity/api/strategies.html). - **`tsn_affinity.cli`** — Command-line entry points. See [Operations → CLI](/tsn-affinity/operations/cli.html).

Top-level re-exports

The following symbols are re-exported from the top-level tsn_affinity namespace. Importing from the top level avoids the cost of importing submodules on hot paths.

=== “Version and core”

```python
from tsn_affinity import __version__
from tsn_affinity import (
    DecisionTransformer,
    DTBackbone,
    ObsEncoder,
)
```

=== “Configs and exceptions”

```python
from tsn_affinity import (
    ModelConfig,
    SparseConfig,
    RoutingConfig,
    TSNAffinityConfig,
    SparseConversionConfig,
    TSNAffinityError,
    RoutingError,
    MaskError,
    ConfigurationError,
    DataError,
    StrategyError,
    BenchmarkError,
    setup_logging,
)
```

=== “Data and routing”

```python
from tsn_affinity import (
    Trajectory,
    discount_cumsum,
    compute_action_affinity,
    compute_hybrid_affinity,
    compute_latent_affinity,
    compute_latent_affinity_batch,
    normalize_scores,
    MaskWarmstarter,
)
```

=== “Sparse layers and strategies”

```python
from tsn_affinity import (
    TSNSparseLinear,
    TSNSparseConv2d,
    TSNSparseEmbedding,
    TopKMaskSTE,
    convert_to_sparse,
    iter_sparse_modules,
    kmeans_quantize,
    rebuild_optimizer,
    BaseStrategy,
    TSNBaseStrategy,
    TSNCoreStrategy,
    TSNAffinityStrategy,
    TSNReplayKLStrategy,
    CopyManager,
    ModelCopy,
)
```

Per-package references