Metadata-Version: 2.4
Name: hybridlens
Version: 0.1.0a1
Summary: Mechanistic interpretability for hybrid Mamba-Attention LLMs (Jamba, Zamba2): cross-component SAEs, ablation attribution, and a Hub registry.
Project-URL: Homepage, https://github.com/hinanohart/hybridlens
Project-URL: Issues, https://github.com/hinanohart/hybridlens/issues
Project-URL: Changelog, https://github.com/hinanohart/hybridlens/blob/main/CHANGELOG.md
Author: hinanohart
License-Expression: Apache-2.0
License-File: LICENSE
Keywords: attention,hybrid-llm,interpretability,jamba,mamba,mechanistic-interpretability,sparse-autoencoder,zamba2
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: Apache Software License
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.11
Requires-Dist: datasets>=2.20
Requires-Dist: huggingface-hub<1.0,>=0.26
Requires-Dist: numpy>=1.26
Requires-Dist: pydantic>=2.7
Requires-Dist: recurrentlens<0.2,>=0.1.0.post1
Requires-Dist: safetensors>=0.4
Requires-Dist: torch<3.0,>=2.4
Requires-Dist: transformers<6.0,>=4.46
Requires-Dist: zarr<3.0,>=2.18
Provides-Extra: attribution
Requires-Dist: captum>=0.7; extra == 'attribution'
Provides-Extra: dev
Requires-Dist: hypothesis>=6.100; extra == 'dev'
Requires-Dist: pyright>=1.1; extra == 'dev'
Requires-Dist: pytest-cov>=5.0; extra == 'dev'
Requires-Dist: pytest>=8.0; extra == 'dev'
Requires-Dist: ruff>=0.6; extra == 'dev'
Provides-Extra: mamba
Requires-Dist: causal-conv1d>=1.4; extra == 'mamba'
Requires-Dist: mamba-ssm>=2.2.2; extra == 'mamba'
Provides-Extra: nnsight
Requires-Dist: nnsight>=0.3; extra == 'nnsight'
Description-Content-Type: text/markdown

# hybridlens

![Alpha](https://img.shields.io/badge/status-alpha-orange)
[![License: Apache-2.0](https://img.shields.io/badge/License-Apache_2.0-blue.svg)](LICENSE)
[![PyPI](https://img.shields.io/pypi/v/hybridlens.svg)](https://pypi.org/project/hybridlens/)

Mechanistic interpretability for **hybrid Mamba-Attention LLMs** —
[Jamba](https://huggingface.co/ai21labs), [Zamba2](https://huggingface.co/Zyphra) —
with cross-component SAEs, ablation attribution, and a Hub registry.

`hybridlens` is the [`recurrentlens`](https://github.com/hinanohart/recurrentlens)
sibling for **mixed-architecture** models, where attention layers, SSM layers,
and MoE routers coexist in a single forward pass. Pure-transformer tools
(SAELens, OpenMOSS) and pure-SSM tools (recurrentlens) both assume a single
component family; hybrid models need per-component analysis.

## Why hybrid-specific tooling

2026's mainline open-weight LLMs are **no longer pure transformers**: AI21's
Jamba 1.5, Zyphra's Zamba2, NVIDIA's Nemotron Nano 2, and Microsoft's Samba
all ship Mamba+Attention+MoE in a single stack. The existing mechanistic
interpretability toolchain (`SAELens`, `OpenMOSS`, `circuit-tracer`) was
written against decoder-only transformers and silently mis-attributes
features when an SSM layer or MoE router is present. `hybridlens` fills
that gap with per-component SAEs trained on the **actual** hidden stream
each layer carries, plus attribution that respects the component boundary.

> **Solo-maintained OSS. Best-effort response. No production SLA.**
> Released 2026 by [@hinanohart](https://github.com/hinanohart) under Apache-2.0.

## v0.1.0 scope

Eight public APIs covering Jamba + Zamba2:

- `load_hybrid_model` — load and tag attention / SSM / MoE components
- `HybridLensModel` — wrapper that **composes** (does not subclass) the
  `recurrentlens` primitives (`BaseSAE`, `HookManager`, `ActivationCache`)
- `LayerLayout` — pydantic schema for per-layer component map
- `train_component_sae` — train an SAE on one component's hidden stream
- `ComponentSAEBundle` — save/load a per-component SAE collection
- `attribute(method='ablation')` — zero-ablation attribution to target metric
- `steer_component` — context-managed activation steering
- `hub_push_bundle` / `hub_load_bundle` — HF Hub I/O for bundles

Deferred to v0.1.1+: captum integrated-gradients attribution, MoE expert
routing trace, Sankey cross-component flow viz, Nemotron / Samba adapters,
pretrained Hub bundles, HF Space deploy.

### v0.1.0a0 alpha — what runs vs what raises

The alpha freezes the eight public signatures and ships the bookkeeping
half of the library — anything that does *not* need a live forward pass
through a hybrid LLM. Concretely:

| API | v0.1.0a0 status |
|-----|-----------------|
| `LayerLayout` / `LayerSlot` | runs |
| `JambaAdapter.discover_layout` / `Zamba2Adapter.discover_layout` | runs (config-driven) |
| `hooks.resolve_site` | runs |
| `ComponentSAEBundle` (in-memory + `save` / `load` round-trip) | runs |
| `viz.render_hybrid_explorer` | runs |
| `hub_push_bundle` / `hub_load_bundle` | runs (lazy HfApi import) |
| `load_hybrid_model` | **`NotImplementedError`** until v0.1.0 final |
| `train_component_sae` | **`NotImplementedError`** until v0.1.0 final |
| `attribute(method='ablation')` | **`NotImplementedError`** until v0.1.0 final |
| `steer_component` (scale ≠ 0 branch) | **`NotImplementedError`** until v0.1.0 final |
| `hooks.install_capture_hook` | **`NotImplementedError`** until v0.1.0 final |

The five remaining APIs require `torch` + `transformers` +
`recurrentlens` to be installed in the same environment; they land in
v0.1.0 (final) after the Phase 5 release-verify venv exercises them
end-to-end. The signatures are frozen — you can write code against them
today.

## Install

```bash
pip install hybridlens
# Optional Mamba kernels (CUDA only):
pip install "hybridlens[mamba]"
```

## Quick look

```python
import hybridlens as hl

model = hl.load_hybrid_model("ai21labs/Jamba-tiny-dev", family="jamba")
bundle = hl.train_component_sae(
    model,
    component="attn",
    layer=4,
    n_tokens=200_000,
)
print(bundle.summary())
```

> ⚠️ **HF Hub anonymous rate limit (100 req/h)** —
> set `HF_TOKEN` for steady-state usage.

## Known limitations (v0.1.0)

- Only Jamba and Zamba2 adapters. Nemotron-Nano / Samba arrive in v0.1.1.
- Attribution is zero-ablation only; captum IG is v0.1.1.
- MoE expert tracing and cross-component Sankey viz are v0.2.
- Default training dataset is Pile-uncopyrighted; users supplying their own
  corpora are responsible for license compliance.
- Bundle Hub assets are user-uploaded; no pretrained release ships with v0.1.0.

## Citation

```bibtex
@software{hinanohart_hybridlens_2026,
  author  = {hinanohart},
  title   = {hybridlens: cross-component SAE attribution for hybrid LLMs},
  year    = {2026},
  url     = {https://github.com/hinanohart/hybridlens},
  version = {0.1.0a0},
}
```

## Related

- [`recurrentlens`](https://github.com/hinanohart/recurrentlens) — sibling
  project for pure-SSM models. `hybridlens` reuses its `BaseSAE`,
  `HookManager`, `ActivationCache`, and `hub.io` primitives.
