[models] Add TFMTokenizer - #1202

CI blocker follow-up on the current head:

• Added the required Foundation Model + Attention/Transformer doc badges.
• The default forward(X) is now Tensor-valued (reconstruction), while the complete self-supervised training payload is exposed explicitly through tokenize(...).
• Removed the private _compute_spectrogram / _encode method chain from the scripted path so convert_model_to_plain() retains every method used by forward.
• Applied the two docstrfmt changes reported by pre-commit.
• Model-specific structured-output tests now call tokenize(...); a regression checks the default Tensor forward contract.

These changes address the four failures from the last upstream run without adding integration-test skips. Could the workflow be re-run on the current head when convenient?

Reference-fidelity correction on current head 5e6d18ef after re-auditing the released TFM-Tokenizer code:

• The released tokenizer updates quantizer.embedding.weight through EMA; after nearest-neighbor lookup it applies quant_out = quant_in + (quant_out - quant_in).detach(), so the subsequent VQ loss does not send optimizer gradients into the codebook table.
• The port now makes that contract explicit: lookup uses a detached codebook, encoder gradients still flow through the straight-through estimator, and the codebook remains EMA-only.
• Regressions now require encoder-path gradients, no codebook optimizer gradient, a real EMA update in train mode, and no second codebook change from the optimizer step.
• I also fixed a misplaced pytest.mark.parametrize decorator in the architecture-validation tests that would otherwise break collection when upstream CI is approved.

No architecture dimensions or reconstruction path changed; this is a correctness/fidelity cleanup before maintainer review.

EMA robustness follow-up on the current head fd494c01:

• I re-audited the released TFM-Tokenizer EMA update and found a library-level edge case inherited from the reference implementation: cluster_size starts at zero while ema_weight starts from the random codebook. After the first training batch, never-selected codes are normalized by eps; with the reference defaults (codebook_size=8192, eps=1e-5) this can inflate an unseen code by roughly 1/eps and move it far outside the L2-normalized encoder space.
• The port now leaves never-selected code vectors at their initialized locations until their first assignment. EMA statistics are still updated exactly as before, and occupied codes still use the same EMA centroid normalization.
• A focused regression uses more codebook entries than assignments and requires all unselected entries to remain unchanged, finite, and near their initialized scale.

This deliberately improves numerical robustness rather than claiming checkpoint parity with the research repository; the tokenizer architecture, VQ objective, straight-through estimator, and checkpoint state schema are unchanged.

VQ gradient audit on the current head: with the straight-through tensor used in both VQ terms, the EMA-only port sent equal-and-opposite gradients into the encoder when commitment_cost=1.0 (the reference default). I changed only the gradient routing, not the objective value: the codebook-distance term now uses the raw detached EMA codebook vector, the commitment term trains the encoder, reconstruction still uses the straight-through path, and the codebook remains EMA-only. The regression now requires a strictly non-zero encoder gradient rather than only checking grad is not None.

One more VQ-gradient invariant is now explicit on head e9a8138a.

The released code applies the straight-through estimator before evaluating both VQ MSE terms. If
q_st = z + (q - z).detach(), then the two encoder gradients are opposite:

• ∇z MSE(q_st, z.detach()) ∝ +(q - z)
• ∇z MSE(q_st.detach(), z) ∝ -(q - z)

At the released/default commitment_cost=1.0, those contributions cancel exactly. That makes the nominal VQ objective produce zero encoder gradient even though the codebook itself is EMA-updated.

The current port keeps the same scalar objective value but evaluates the codebook term from the detached EMA codebook vector itself. The codebook term is therefore value-only (the codebook stays EMA-only), while the commitment term supplies the encoder gradient. A focused regression now requires nonzero frequency- and temporal-encoder gradients at commitment_cost=1.0 and still requires no optimizer gradient on the codebook embedding.

This is a deliberate library-level correctness fix rather than a checkpoint-parity claim; reconstruction architecture and EMA state remain unchanged.

Pinned official-checkpoint parity is now available for the current PR head abb99259f8da6a02e0c0366e7feb39363e3a617c.

Reference: Jathurshan0330/TFM-Tokenizer@2d6da482b16dabbb2ebec808fa9f505fc6f367c4; official 2x2x8 checkpoint SHA256 f7929eed0f8275d08a7e18743fdd188c7ba5b4722151333b4d386ca05b84d6d9.

CPU assay results:

• all 191 checkpoint state entries mapped;
• token agreement: 1.000;
• spectrogram / pre-quant embeddings / straight-through quantized embeddings / reconstruction: max abs error 0;
• reconstruction-path gradients for raw EEG input, first frequency conv, and first temporal conv: max abs error 0.

Run: https://github.com/lindicaphxag-tech/lindicaphxag-tech/actions/runs/37114458814

The EMA-VQ objective-gradient correction remains an explicitly documented library-level deviation and is intentionally excluded from the reference-parity oracle rather than presented as reference behavior.

@bruAristimunha @julien-gadonneix — this is now at the point where I would value a maintainer pass on the model/API contract rather than further self-hardening. The current head keeps the default forward(X) Tensor-valued, isolates the full SSL payload in tokenize(...), makes the EMA-only codebook gradient route explicit, and includes checkpoint/config/EMA regressions. I’ll avoid additional scope unless review identifies a concrete issue.

Follow-up on the EMA regression: the failure exposed a real update bug rather than a flaky test. embedding.weight[occupied].copy_(...) was writing into the temporary tensor produced by boolean advanced indexing, so selected EMA centroids never reached the actual codebook. Head eaf26ab1 switches that one line to indexed assignment; unseen entries still remain unchanged. Independent current-head validation is green: https://github.com/lindicaphxag-tech/lindicaphxag-tech/actions/runs/37116304219. I’ll hold the PR here for maintainer review.

Review-history cleanup only: I squashed the branch to a single commit before maintainer review. The Git tree is unchanged (c6e5121cf25a8f4c2bb91ac5ed18fd87822c8a61), so the existing focused/parity evidence remains applicable. I’ll hold this head unless review identifies a concrete issue.

Codecov Report

❌ Patch coverage is 92.07921% with 8 lines in your changes missing coverage. Please review.
✅ Project coverage is 89.19%. Comparing base (83a1e67) to head (37afb53).
⚠️ Report is 2 commits behind head on master.

@@            Coverage Diff             @@
##           master    #1202      +/-   ##
==========================================
+ Coverage   89.16%   89.19%   +0.02%     
==========================================
  Files         158      159       +1     
  Lines       19641    19741     +100     
==========================================
+ Hits        17513    17607      +94     
- Misses       2128     2134       +6