Add AST (Audio Spectrogram Transformer) Adapter

[dylanberens]

Description

This PR introduces full support for the Audio Spectrogram Transformer (AST) architecture with TransformerLens.

Key Engineering Details
Because AST is built on a ViT-base backbone but processes continuous audio spectrogrmas, several important architectural bridges were necessary:

• Custom ASTEmbed: Implemented a specialized Conv2d embedding layer to handle 1D to 2D tensor reshaping and extract overlapping patches using a stride=(10, 10), which is a frequent cause of disconnects whens working with AST attention weights
• Dual-Token Pooling: Implemented the AST-specific MLP head pooling logic, which averages botht eh CLS and Distillation tokens before classification.
• Dual LayerNorm: Accounted for AST's second LayerNorm applied exclusively to the pooled token prior to the embedding matrix

Adds ast.py (architecture adapter) and test_ast_adapter.py (verifying FP32 mathematical parity with max logit diff < 1e-4).

Fixes # 1484

Type of change

• New feature (non-breaking change which adds functionality)

Checklist:

• I have commented my code, particularly in hard-to-understand areas
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes
• I have not rewritten tests relating to key interfaces which would affect backward compatibility

[jlarson4]

Thanks for putting this together @dylanberens! The AST domain work is great, the Q/K/V/O and bias reshapes, MLP transposes, the overlapping-patch Conv2d, the cls-distillation-patches token order, and the pooled dual-token + dual-LayerNorm head are all correct.

We will need to address & adjust the approach: this is written against the legacy HookedTransformer convert_weights path, but TransformerBridge adapters are driven by component_mapping (wrap HF's live modules). As submitted it can't load or run, which is why your tests are failing in CI.

Primary issues with the current code:

1. Doesn't integrate with the bridge – ASTAdapter never sets self.component_mapping (it defines no __init__, so it inherits None), and it isn't registered in the factory. Booting AST raises ValueError: Unsupported architecture. get_config_map/convert_weights are legacy-HT classmethods no bridge path ever calls.
2. The test fails and never touches the bridge – It KeyErrors at ast.py:125, the HF keys are wrong on two axes: layers.{i} should be encoder.layer.{i}, and q_proj/k_proj/v_proj/o_proj/mlp.fc1/fc2 should be attention.attention.{query,key,value} / attention.output.dense / intermediate.dense / output.dense. So the max diff < 1e-4 claim isn't reproducible, and the test builds a raw HookedTransformer, not a bridge.

To make this mergeable:

• Rewrite as a component_mapping adapter: I'd recommend following our Architecture Adapter Creation Guide. The hubert.py adapter is our current working audio adapter, it would be a good reference point to work off.
• Delete ASTEmbed and both classmethods — the bridge wraps HF's own modules, so none of that math needs reimplementing.
• Fix n_ctx = num_patches+2; read patch/stride/hidden_act from config.
• Register ASTForAudioClassification in the factory + export it from __init__.
• Rewrite the test to boot a real TransformerBridge and assert bridge-vs-HF fp32 parity, under tests/unit/....

Also, a couple repo items: make sure to drop poetry.lock from the PR. Starting from 3.0.0, the repo is uv-based. And, base this change to dev, we will be able to launch this as part of a 3.x release, it does not need to be gated behind version 4.

Happy to re-review once it's on the TransformerBridge pattern.