Add NVIDIA RE-USE (SEMamba) speech enhancement model - #551
Open
justinchuby wants to merge 13 commits into
Open
Conversation
RE-USE is NVIDIA's universal speech enhancement generator: a dense convolutional encoder, 30 time-frequency bidirectional Mamba blocks, and separate magnitude/phase decoders. It ships a bespoke config.json with no model_type and a mamba_ssm-based checkpoint, so neither transformers AutoConfig nor the existing decode-time Mamba components applied. Components: - SequenceSelectiveScan: the Mamba1 (S6) recurrence over a whole sequence with a zero initial state, expressed as a single ONNX Scan. dA and dBx are computed inside the scan body instead of being materialised for all timesteps, which keeps the working set at (batch, seq, d_inner) rather than (batch, seq, d_inner, d_state). A rides along as a pass-through carry so the body never reaches into the enclosing graph's scope. The recurrence stays in float32, matching the reference kernel. - SequenceMambaBlock: the stateless full-sequence counterpart to MambaBlock, sharing its parameter names and shapes. - Conv2d gains an optional trailing `dilation` argument (needed by the dense blocks); existing positional calls are unaffected. Model and task: - models/reuse.py with ReUseConfig (parsed from the published model_cfg / stft_cfg sections), the encoder/decoder stack, bidirectional Mamba, and a hand-rolled atan2 since ONNX has no Atan2 op. - SpeechEnhancementTask: noisy_mag/noisy_pha in, denoised_mag/pha/com out. The STFT and ISTFT stay outside the graph, as for every other audio model here. - build_reuse() loads config + weights from a directory or the Hub, since build() cannot discover a repo with no model_type. Module attribute names are structured so encoder/decoder parameters land on the checkpoint's nn.Sequential indices directly; preprocess_weights only has to nest the flat SSM parameters under `ssm`, the same offset the Mamba causal-LM models correct for. Verified against a pure-PyTorch transcription of the reference model with the real nvidia/RE-USE weights: magnitude and complex outputs match to 4.2e-06, and every initializer is populated. The phase output matches to 1.8e-04 at 10 of 6601 points, all of which are the smallest-radius points where atan2 is ill-conditioned; the error in the complex plane there is below 1.7e-05. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 43bb31ac-c136-44bc-8cdf-70cbcedd4a42 Signed-off-by: Copilot <justinchuby@users.noreply.github.com>
Benchmarking the exported model against the MLX plugin EP
(onnxruntime-ep-mlx 0.29.1, Apple M1 Max) showed no speedup at all: MLX
claimed 318 of 2567 nodes as 318 single-node partitions, and every one of
the 120 Scan recurrences fell back to CPU.
The EP's own diagnostics (ONNXRUNTIME_EP_MLX_VERBOSE=1) named the reason:
Scan x120: scan input 2 must have a statically known non-empty axis 0,
got shape [-1, -1, 256]
`scan_input_axes=[1, ...]` avoided a pre/post transpose but put the scan
axis where the EP will not take it. Feeding the Scan time-major instead
lets it claim all 120 recurrences, fusing 3533 of 5333 nodes into 121
compiled subgraphs. This matches what functions/linear_attention.py
already does, and the transposes cost far less than the fallback.
Measured on nvidia/RE-USE with static input shapes, median of 5 runs:
101 frames (0.51s audio) CPU 1392 ms -> MLX 512 ms 2.72x
401 frames (2.00s audio) CPU 4431 ms -> MLX 1346 ms 3.29x
At 2s the model goes from RTF 2.21 to RTF 0.67, i.e. from well under to
comfortably over real time. Outputs stay within 7.2e-04 of the CPU run.
The scan length is only static when the graph inputs are, so this helps
fixed-size/chunked export; a fully dynamic model still runs the
recurrence on CPU.
Also documents, and pins with a test, why the backward branch reverses
with a negative-step Slice. Switching it to ReverseSequence looks like a
strict improvement -- the EP claims it and it is ~4x faster in isolation
-- but the unclaimable Slice is what splits the model into one small
compilable subgraph per block. Claiming the reverse collapsed all 30
blocks into a single subgraph that the EP ran eagerly (eager=7, no
compile cache), costing 70s instead of 0.5s.
Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
Copilot-Session: 43bb31ac-c136-44bc-8cdf-70cbcedd4a42
Signed-off-by: Copilot <justinchuby@users.noreply.github.com>
`write_onnx_genai_config` had no detector for a RE-USE package, so it fell through to the decoder branch and failed with "decoder workflow requires rank-2 token input and rank-3 logits output". An enhancement model is not generative: it reads a whole spectrogram at once, carries no state between calls and has no logits to sample, so decoder metadata would have published a generation loop the artifact cannot execute. Adds a structural detector (single `model` consuming noisy_mag/noisy_pha and emitting denoised_mag/denoised_pha) plus a builder that emits a single pure invocation, modelled on the existing encoder-embedding writer. The graph consumes a spectrum rather than a waveform, so when the config carries the STFT geometry the document also declares an audio preprocessing program -- decode, resample, downmix, spectrogram, log1p -- and the workflow accepts encoded audio through the same audio-preprocess adapter ABI the CTC path uses. A caller that computes the STFT with a different window or hop, or that skips the log1p compression RE-USE trains on, would silently feed the model out-of-distribution input; publishing the program makes the contract self-describing. When the geometry is absent nothing is invented: the preprocessing section is omitted and the caller supplies the spectra directly. Two details worth noting: - The magnitude/phase ports use an opaque role in the geometry-less path. The portable runtime-role vocabulary has no term for a spectrogram, and labelling them as audio samples would be wrong. - The audio output bindings publish a full `contract`, which the schema requires whenever the package declares a `pipeline.workflow`. Verified against onnx-genai's committed inference_metadata.schema.json, both with and without the preprocessing program, and end to end on the real nvidia/RE-USE checkpoint: the emitted STFT geometry matches the published config (8 kHz, n_fft 320, hop 40, win 320) and the declared artifact loads with exactly the declared ports. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 43bb31ac-c136-44bc-8cdf-70cbcedd4a42 Signed-off-by: Copilot <justinchuby@users.noreply.github.com>
Performance Comparison
|
![github-code-quality[bot]](https://avatars.githubusercontent.com/u/9919?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Pull request overview
Adds NVIDIA’s RE-USE (SEMamba) spectral speech-enhancement model to Mobius, including new full-sequence Mamba/SSM components, a dedicated speech-enhancement task/IO contract, and onnx-genai workflow metadata emission + tests.
Changes:
- Introduces full-sequence (stateless) Mamba-1 selective scan via ONNX
Scan(SequenceSelectiveScan) and a matchingSequenceMambaBlock. - Adds the RE-USE model implementation + bespoke
SpeechEnhancementTask(noisy mag/phase → denoised mag/phase/complex). - Extends onnx-genai auto-export + workflow metadata to detect and describe speech-enhancement packages, with schema-validation tests.
Reviewed changes
Copilot reviewed 18 out of 18 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/build_graph_test.py | Excludes RE-USE/SEMamba from generic config matrices (dedicated tests cover it). |
| src/mobius/tasks/_speech_enhancement.py | New task defining spectral enhancement IO contract and publishing output shapes. |
| src/mobius/tasks/init.py | Exports/registers SpeechEnhancementTask. |
| src/mobius/models/reuse.py | Implements RE-USE model, config parsing, bidirectional TF-Mamba blocks, and build_reuse(). |
| src/mobius/models/reuse_test.py | Model/unit/runtime/EP invariants tests (including atan2 correctness and partitioning invariants). |
| src/mobius/models/init.py | Exposes RE-USE model/config/build entry points. |
| src/mobius/integrations/onnx_genai/workflow_metadata.py | Adds speech-enhancement workflow metadata builder + writer and STFT preprocessing program emission. |
| src/mobius/integrations/onnx_genai/speech_enhancement_metadata_test.py | Validates detection/metadata/schema behavior for enhancement packages. |
| src/mobius/integrations/onnx_genai/auto_export.py | Adds structural detection + dispatch to speech-enhancement metadata writer. |
| src/mobius/integrations/onnx_genai/init.py | Exports new metadata builder/writer APIs. |
| src/mobius/components/_ssm.py | Adds SequenceSelectiveScan implemented as a single ONNX Scan over the sequence. |
| src/mobius/components/_ssm_test.py | Unit tests for SequenceSelectiveScan graph structure/precision/axes invariants. |
| src/mobius/components/_mamba_block.py | Adds SequenceMambaBlock using SequenceSelectiveScan for offline/full-sequence use. |
| src/mobius/components/_mamba_block_test.py | Tests parameter parity and numerical parity vs a PyTorch reference recurrence. |
| src/mobius/components/_conv.py | Adds dilation support to Conv2d. |
| src/mobius/components/_conv_test.py | Tests dilation behavior and preserves positional-call ABI. |
| src/mobius/components/init.py | Exports new SSM/Mamba sequence components. |
| src/mobius/_registry.py | Registers reuse/semamba model_type fallbacks to the new model/task/config. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
🏗️ Architecture Diff
No architecture changes detected. ✅ Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed) |
justinchuby
added a commit
that referenced
this pull request
Aug 23, 2026
… widths (#552) ## The bug `_MimiConvEncoder` in `src/mobius/models/qwen3_tts_tokenizer.py` took **no config at all** and hardcoded its channel widths and kernels (`1→64→128→256→512→1024→512`). Its output fed `CodecEncoderTransformerModel(hidden_size=config.codec_encoder.hidden_size)`, which has no input projection — so the conv output width and `hidden_size` had to agree, and nothing checked that they did. When they disagreed, the builder happily emitted a malformed graph: ``` Conv encoder/encoder/layers/14/conv weight [512,1024,3] -> [batch, T, 512] LayerNormalization weight [32] bias [32] over a 512-wide tensor <-- malformed MatMul q_proj [32,32] vs 512 -> "Incompatible MatMul contraction dimensions: 512 vs 32" ``` ## Why it was invisible `SymbolicShapeInferencePass` catches inference errors and logs `"Symbolic shape inference failed (upstream bug); skipping"` — a deliberate generic guard. The tiny config in `TestBuildCodecGraph._codec_config()` set `codec_encoder.hidden_size=32`, so **the codec tests were building a broken graph and passing**. ## The fix The hardcoded widths were not dictated by the checkpoint — HF's `MimiEncoder.__init__` derives the entire stack from config, and mobius's literals were just that derivation inlined with `num_filters=64`, `upsampling_ratios=[8,6,5,4]`, `num_residual_layers=1`, `kernel_size=7`, `last_kernel_size=3`. So rather than adding a validation guard, this parameterizes the stack and eliminates the mismatch class: - **`CodecEncoderConfig`**: added the missing HF fields (`audio_channels`, `num_filters`, `num_residual_layers`, `kernel_size`, `last_kernel_size`, `residual_kernel_size`, `compress`, `upsampling_ratios`) with HF's own defaults, extracted in `ArchitectureConfig.from_transformers`. - **`_MimiConvEncoder`**: builds by HF's derivation. The parameterless ELU slots are preserved so `encoder.encoder.layers.*` numbering still matches the checkpoint exactly. `kernel = ratio * 2` and the trailing conv's out-channels is `hidden_size`, both derived. - **`_EncoderResBlock`**: inner width is now `dim // compress` with a configurable kernel. ## Also fixed: the encoder RVQ could not load the real checkpoint This one is a **weight-loading correctness fix, not just a shape-consistency fix**, which is why it belongs here despite being adjacent to the original scope. `_EncoderSplitRVQ` used `input_dim=codebook_dim` and `dim=codebook_dim // 2`. The real checkpoint has `encoder.quantizer.*.input_proj.weight [256, 512, 1]` and `codebook.embed_sum [2048, 256]`, with `encoder_config.codebook_dim=256`. So the old code built `Conv1d(256→128)` projections and 128-wide codebooks: **the encoder could never have loaded the real checkpoint's weights**, independently of the malformed-graph issue. It only happened to be self-consistent on the `enc is None` fallback path, where `codebook_dim` defaulted to 512. Now it projects `hidden_size → codebook_dim`, matching HF `MimiResidualVectorQuantizer`. ### The encoder/decoder asymmetry is intentional — please don't unify these paths The two sides interpret `codebook_dim` differently, and both are correct: | | config value | RVQ input | codebook dim | checkpoint `input_proj` | |---|---|---|---|---| | encoder | `codebook_dim=256` (Mimi semantics: the codebook dim itself) | `hidden_size=512` | 256 | `[256, 512, 1]` | | decoder | `codebook_dim=512` (the RVQ *input* dim) | 512 | `512 // 2 = 256` | `[256, 512, 1]` | So the decoder's `// 2` convention in `SplitResidualVectorQuantizer` is right and is left untouched; only the encoder path was misreading its field. A future "cleanup" that unifies the two would reintroduce this bug. ## Verification **Symbolic-shape-inference failures across `TestBuildCodecGraph`** (measured by temporarily instrumenting `infer_symbolic_shapes`, since the pass swallows them — instrumentation not committed): | | failures | |---|---| | before | **6** | | after | **0** | **Default config reproduces the real checkpoint exactly** — verified against the safetensors header of `Qwen/Qwen3-TTS-Tokenizer-12Hz`: all 14 conv-stack weights match by name and shape, including the ResBlock `layers.N.block.1/.3` names and the `0,1,3,4,6,7,9,10,12,14` numbering, as do both RVQ `input_proj`/`output_proj` pairs and the codebooks. This is asserted by tests, so indexing or width drift (which would otherwise surface far away at weight-load time) fails loudly. **Tiny test config** keeps the real depth and only shrinks widths (`num_filters=4`, `hidden_size=32` → `1→4→8→16→32→64→32`), so it exercises the same weight-name structure the real model depends on. Codec test runtime: 2.05s (4 tests) → ~2.1s (9 tests). **Suite**: `pytest tests/build_graph_test.py tests/cli_test.py src/ -k "not phi4mm and not apply_weights_unknown" -n auto` → 8 failed, 4613 passed — the same 8 pre-existing failures on `main` (onnx_genai jsonschema drift + one qwen_image golden). `lintrunner` clean. ## New tests - `test_default_config_conv_stack_matches_checkpoint` — asserts derived parameter **names** and shapes against the real checkpoint. - `test_tiny_config_keeps_checkpoint_layer_names` — tiny config has identical naming, narrower tensors. - `test_conv_stack_indices_follow_config` — 2 ratios → final conv at `layers.8`; 4 ratios × 2 residual layers → `layers.18`. - `test_hidden_size_drives_conv_output_width` — direct regression guard for the original bug. - `test_default_config_rvq_projections_match_checkpoint` — asserts `input_proj (256, 512, 1)`, `output_proj (512, 256, 1)` and `(2048, 256)` codebooks, guarding the weight-loading fix. ## Decoder conv stack Checked: the codec **decoder** does **not** have the mirrored conv problem. `Qwen3TTSCodecDecoderModel` already derives everything from config (`latent_dim`, `decoder_dim`, `upsample_rates`, `upsampling_ratios`), and `CodecDecoderTransformerModel` has an `input_proj`/`output_proj` pair that adapts `latent_dim ↔ hidden_size`, so no width can be stranded. No change needed there. ## Scope Limited to the codec tokenizer model, its config, and its tests. Does not touch `reuse.py`, `_ssm.py`, `_mamba_block.py`, the onnx_genai integration (covered by #551), or `SymbolicShapeInferencePass`. --------- Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com> Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
tests/model_coverage_test.py enforces that every registered model_type is
either driven by a config in tests/_test_configs.py or carries a documented
reason in _COVERAGE_SKIP. Registering "reuse" and "semamba" without either
broke six of its checks. I had only been running the command in the repo
instructions (build_graph_test, cli_test, src/), so this file and
tests/weight_alignment_test.py were never exercised.
Rather than skipping the coverage requirement, drive the model from it:
- Add a "reuse" entry to SPEECH_CONFIGS using the `_config_cls` escape hatch
that ParakeetCTCConfig already uses, so TestBuildSpeechGraph builds it for
real (package shape, initializers, ONNX checker, output shapes/dtypes).
- Declare `speech-enhancement -> {"model"}` in _SPEECH_TASK_KEYS so that
harness asserts the package shape instead of silently accepting any keys.
- _COVERAGE_SKIP records only what is genuinely unavailable: L2 for "reuse",
because arch_validation_test asserts the downloaded config has a
model_type and the published nvidia/RE-USE config.json is a bespoke
model_cfg/stft_cfg document without one; and "semamba" as a bare alias.
The per-model L1/L3 check passes on the config regardless of the skip
entry, so this documents the L2 gap without waiving L1/L3.
That new coverage immediately caught a real bug. weight_alignment_test
requires preprocess_weights to be an identity on already-correct names, and
mine renamed on suffix alone: given `forward_blocks.ssm.A_log` it produced
`forward_blocks.ssm.ssm.A_log`, dropping all 20 SSM parameters. Loading an
already-converted state dict would have silently lost every selective-scan
weight. Fixed by skipping keys already nested under `ssm`, with a test
asserting a second pass is a no-op.
Verified against the branch tip with the same full sweep (tests/ + src/,
-m "not integration"): 21 failures -> 15, the six removed being exactly the
coverage failures, and zero new. None of the remaining 15 touch this work.
Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
Copilot-Session: 43bb31ac-c136-44bc-8cdf-70cbcedd4a42
Signed-off-by: Copilot <justinchuby@users.noreply.github.com>
The comment and test justified the negative-step Slice on the grounds that
the MLX plugin EP cannot claim it, so it acts as a partition boundary that
keeps each block a separately compilable subgraph -- with the claim that
switching to ReverseSequence cost 70s instead of 0.83s end to end.
That measurement was taken against onnxruntime-ep-mlx 0.29.1. Since 0.29.2
("Compile ONNX control flow with MLX") the EP compiles If/Loop/Scan inside
the surrounding closure instead of forcing the fused subgraph eager, and it
implements negative-step Slice. Both halves of the old reasoning are gone:
the reverse Slice is now claimed, and there is no eager collapse to avoid.
Keep the Slice -- it is still the better spelling, for reasons that do not
depend on any EP's coverage: it is the standard way to express a reversal,
whereas ReverseSequence needs an extra Expand to build per-row lengths and
implies padding semantics this model does not have, since nothing here is
padded and every row is reversed in full.
Only the justification changes; the emitted graph is identical.
Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
Copilot-Session: 43bb31ac-c136-44bc-8cdf-70cbcedd4a42
Signed-off-by: Copilot <justinchuby@users.noreply.github.com>
justinchuby
added a commit
that referenced
this pull request
Aug 23, 2026
…ive contract (#554) Fixes six long-standing schema-conformance failures caused by mobius's legacy onnx-genai emitters drifting from the runtime's published JSON schema. ## Why the drift was invisible The conformance tests searched a couple of hard-coded local onnx-genai checkout paths and `pytest.skip`ped when none was found. CI has no such checkout, so **six tests never ran there**; on the machines that did have one, the result depended on whatever revision that clone happened to sit on. Two upstream contract redesigns accumulated unnoticed. ## The two contract redesigns 1. **`pipeline` is now `PipelineSpec`** — a single property, `workflow` (a typed SSA graph), with `additionalProperties: false`. The legacy emitters produced `{models, dataflow, strategy, phases}`, all of which the schema now rejects. 2. **`speculative` is now `SpeculativeContract`** — requires `{proposer, target, vocabulary, max_proposal_width}` and forbids every field of the old flat MTP block. The flat `SpeculatorConfig` that block was modelled on still exists upstream, but it describes a HuggingFace `config.json` section, not `InferenceMetadata.speculative`. ## What was migrated ### Diffusion — `build_diffusion_pipeline_metadata` Now emits the denoise loop explicitly: the solver, sigma/alpha schedule, timestep table, guidance combine and output clamp are real ONNX components built from mobius's policy library and shipped with the document. It **cannot** delegate to `build_diffusion_workflow_metadata`: the ComfyUI conversion path has no component graphs at all (it deliberately does not build or export them), and that builder derives everything from live `ir.Model`s. So the workflow is built directly, reusing `workflow_metadata`'s `_invoke` / `_publish_workflow_v1` and `mobius.generation`'s solver builders. The schedule is derived from the scheduler's own betas through the same diffusers-compatible helpers the package exporter uses, so a ComfyUI conversion and a package export of the same checkpoint describe the same dynamics. img2img's `start_step` lowers to a sliced schedule; the VAE `scaling_factor` and the sigma-space initial-state scale are emitted as explicit components. Three cases now **fail closed** instead of being silently mis-described: - an ancestral sampler (mobius ships no stochastic solver), - Karras/exponential sigma spacing (the workflow ships the sigma table as a constant, so a `use_karras_sigmas` hint is no longer sufficient), - a latent-only graph with no VAE decode. ### MTP — `write_mtp_speculator_metadata` Emits a `SpeculativeContract` anchored to the backbone's workflow: it registers the head as a workflow component, names the target by its declared `logits` port role, states the hidden handoff as `port_bindings.target_hidden_context` plus a `hidden_states` role on the target output, and completes the rollback capacity its own claim requires. It is declared `block`, not `chained`: a chained proposer must expose a `logits_output` carrying the next-token distribution, and this sidecar emits only `mtp_hidden` — the runtime decodes it through the target's shared LM head, which is why that initializer is listed in `shared_weights`. ### VLM `build_native_vlm_package_metadata` produces mobius's **internal** structural descriptor, not a publishable document — `build_vlm_workflow_metadata` already consumes it and republishes only `preprocessing` under a real `pipeline.workflow`. The actual bug was that `write_native_vlm_package_metadata` wrote that descriptor to `inference_metadata.yaml`. It now writes the workflow document, and the tests validate the published document instead of the descriptor. ### CI visibility The upstream schema is vendored under `src/mobius/integrations/onnx_genai/_schema/` and is the default, so conformance never skips and drift becomes a test failure. A local checkout is no longer consulted implicitly — one that is ahead of or behind `main` reintroduces exactly the machine-dependent result that hid this bug. Set `ONNX_GENAI_SCHEMA` to validate against a specific revision. Three further conformance tests (codec, speech-to-text, duplex workflows) were skipping for the same reason and now run. ## Tests Migrated in the same change rather than loosened: - `TestMtpSpeculatorMetadata::test_exact_schema_keys_and_values` asserts the new contract; `test_no_legacy_field_names` is repointed at the now-legacy field names. - A new MTP test anchors against a workflow `write_decoder_workflow_metadata` actually emits (which ships ~11 generated policy components), so target selection can't regress to "the only ONNX component". - New tests pin that the solver schedule comes from the scheduler's betas rather than a placeholder ramp. - Diffusion/ComfyUI tests assert workflow structure instead of the removed `strategy` block; scheduler facts that were duplicated into the document are asserted on the parsed `ComfyUIWorkflow` where they actually live. ## Verification Rebased onto current `main` (includes #552 and #553). ``` python -m pytest tests/build_graph_test.py tests/cli_test.py src/ -q \ -k "not phi4mm and not apply_weights_unknown" -n auto ``` | | before (`main` at branch point) | after | |---|---|---| | full suite | 8 failed / 4613 passed | **0 failed / 4641 passed** | | `src/mobius/integrations/onnx_genai/` with `ONNX_GENAI_SCHEMA` at upstream `main` | 7 failed / 240 passed | **258 passed / 1 skipped** | Six of those eight failures are the schema-drift ones fixed here. The other two — `qwen_image_test.py::test_deterministic_l4_l5_image_edit_golden` and `TestNativeVlmPackageMetadata::test_cached_gemma_processor_matches_emitted_patch_budget` — were fixed by #553, which is now merged; this branch preserves both of its changes through the rebase. `lintrunner` is clean. ## Sequencing note `workflow_metadata.py` and `auto_export.py` are **not modified** (both are read from only), so this does not conflict with #551. --------- Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com> Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
…erage Three small changes, the last two inherited from the main merge rather than from this branch's work. 1. SequenceSelectiveScan's zero carried state dropped its `value` attribute. ONNX already defines ConstantOfShape's default as a float32 zero, so the attribute only restated the default. 2. tests/model_coverage_test.py failed for glm_moe_dsa, which landed in #555 with a test_model_id but no YAML case. Added one, with a skip reason that states why L4/L5 cannot run: GLM-5.2 is 753B parameters across 282 safetensors shards. 3. glm_moe_dsa had no entry in _test_configs.py, so it fell back to an auto-generated config with qk_nope_head_dim/qk_rope_head_dim unset, and the shared MLA component raised `TypeError: unsupported operand type(s) for +: 'NoneType' and 'NoneType'`. Added a tiny config mirroring the model's own fixture in src/mobius/models/glm_moe_dsa_test.py. That takes glm_moe_dsa from 5 failures to 3. The remaining ones are a real defect in the model, not the config: with a buildable graph it now emits an invalid Squeeze in the MLA attention, and ORT refuses to load the model at all: Node (model/layers.0/self_attn/Squeeze_node_46) Op (Squeeze) [ShapeInferenceError] Dimension of input 1 must be 1 instead of 4 which also explains its onnx_checker failure. Fixing that needs GLM-5.2 DSA knowledge and belongs with #555; filed separately rather than guessed at here. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 43bb31ac-c136-44bc-8cdf-70cbcedd4a42 Signed-off-by: Copilot <justinchuby@users.noreply.github.com>
#560 landed the real fix for glm_moe_dsa while this branch was open: a dedicated GlmMoeDsaTask supplying the packed KV-cache shape the model's _unpack_past expects, plus its own tiny config and a _COVERAGE_SKIP entry. That makes both of my stopgaps redundant, and one actively harmful: * tests/_test_configs.py had two "glm_moe_dsa" entries after the merge, so every parametrized case ran twice as glm_moe_dsa_0 / glm_moe_dsa_1 and each failure was reported twice. Removed mine; #560's is shape-shrunk from the real zai-org/GLM-5.2 config.json and is better sourced. * testdata/cases/causal-lm/glm-5.2.yaml existed only to satisfy model_coverage_test, which #560 now satisfies via _COVERAGE_SKIP. A YAML case that can never run (753B across 282 shards) is worse than an explicit skip with a stated reason. The invalid Squeeze I reported in #557 is fixed by #560 as well, so synthetic_parity[glm_moe_dsa] passes again rather than being skipped. Two glm_moe_dsa failures remain (test_onnx_checker_passes, test_outputs_have_shapes_and_dtypes). Verified those reproduce on a clean origin/main worktree, so they are not from this branch and not from the dedupe. Full suite on this branch: 11 failed / 7169 passed. All 11 pre-existing — 8 arch_validation (network-dependent), the 2 glm_moe_dsa above, and synthetic_parity[granitemoehybrid]. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 43bb31ac-c136-44bc-8cdf-70cbcedd4a42 Signed-off-by: Copilot <justinchuby@users.noreply.github.com>
TFMambaBlock read the frequency extent back with Shape(x)[3] even though it is a
fixed function of n_fft: the input frequency axis is n_fft // 2 + 1, forward
pads its tail by 2, and DenseEncoder's dense_conv_2 is a kernel-3, stride-2
convolution with no padding. Emitting it as a constant says that directly
instead of asking the graph at run time.
The three geometry numbers are hoisted to module constants shared by all three
users (the pad in forward, the Conv2d in DenseEncoder, the derivation in
ReUseConfig.encoder_freq_bins), so those cannot drift apart.
Verified bit-exact: same weights, same input, all three outputs identical to the
previous graph (max|delta| = 0.0, np.array_equal true). Also executed in ORT for
n_fft in {320, 400, 322} x time in {7, 32} — 322 being the n_fft % 4 == 2 case
where the floor division actually truncates. Shape nodes 26 -> 25.
On the motivation, which did not survive measurement:
The suggestion was that this would let an EP claim the frequency-axis scans
under a dynamic time axis. It does not, because there was nothing to fix. With
the MLX plugin EP the frequency scans are already claimed, before and after —
0 CPU-resident Scans either way, identical fused-subgraph and CPU-node counts.
Shape inference already resolves the extent to a literal (81 for n_fft=320)
without help, since Shape(x)[3] reads an axis that is itself statically known;
ORT reports [81, batch*(floor(time/4 + 1/4) + 1), 32] for the scan input on
both versions.
Two things had to be corrected before that measurement meant anything, and both
initially pointed the wrong way:
* The environment had an editable onnx-shape-inference 0.1.10 shadowing this
repo's >=0.3.1 pin, so every build printed "Symbolic shape inference failed
(upstream bug); skipping" and every shape came back unknown. That is a stale
install, not an upstream bug.
* The EP being loaded was the released 0.29.1 wheel, which still declines
ConstantOfShape with a runtime shape input and so pushed every Scan to CPU.
Measured against a build of the current EP instead.
Before both corrections the change looked like it helped (32 -> 30 fused
subgraphs); after them the difference is zero. Keeping it anyway on its own
merits — one less runtime read of a build-time constant, and the extent is now
explicit in the graph rather than contingent on an inference pass running.
TestEncoderFreqBins covers the arithmetic across four n_fft values, the
decoder-coverage invariant (2 * extent >= input width, since forward crops the
overshoot), the extent the built graph actually carries, and that no Shape node
feeds the frequency axis again. Confirmed each guard fails when the thing it
guards is perturbed.
Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
Copilot-Session: 43bb31ac-c136-44bc-8cdf-70cbcedd4a42
Signed-off-by: Copilot <justinchuby@users.noreply.github.com>
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 20 out of 20 changed files in this pull request and generated no new comments.
Suppressed comments (3)
src/mobius/models/reuse.py:171
- ReUseConfig.from_json defaults num_tfmamba to 4 when the field is missing, which is inconsistent with the dataclass default (30) and would silently build a much smaller model for partially-specified configs. The fallback should match the class default (or the reference model).
num_tfmamba=int(model_cfg.get("num_tfmamba", 4)),
src/mobius/models/reuse.py:508
- BiMambaBlock hard-codes LayerNorm eps=1e-5 instead of using the config's norm_epsilon. If a checkpoint/config changes epsilon, this will break numerical parity for the TF Mamba blocks.
self.norm = LayerNorm(d_model, eps=1e-5)
src/mobius/integrations/onnx_genai/workflow_metadata.py:10384
- When the workflow includes the audio_preprocess adapter (i.e., transforms is not None), the workflow manifest should also declare the required adapter ABI in manifest.adapter_abis. Other workflows in this file do this for adapter components (e.g. image_preprocess/audio_preprocess), and omitting it can prevent runtimes from discovering/validating required adapters up front.
workflow = {
"manifest": {
"capabilities": ["workflow_ssa", "linear_effects", "typed_emit"],
},
"effects": effects,
Three of the four review findings were real. Each fix is pinned by a test that was confirmed to fail when the fix is reverted. 1. `from_json` defaulted `num_tfmamba` to 4 while the dataclass declared 30. A config.json missing that key would have silently built a model an eighth of the real depth and still loaded, since every block is independently named. Now 30, and a new test asserts every `from_json` fallback equals the dataclass default rather than checking this one field, so the next added field cannot drift the same way. 2. The BiMamba LayerNorm hardcoded eps=1e-5 while the InstanceNorms took `config.norm_epsilon`. Nothing diverged in practice — 1e-5 is both PyTorch's default and the published config's value — but a config setting a different epsilon would have been half-applied. Threaded through; the test builds with norm_epsilon=3e-3 and asserts every normalization node in the graph carries it. Worth noting the reference implementation reads `norm_epsilon` nowhere at all: `mamba_block2_SEMamba.py` uses `nn.LayerNorm(d_model)` and `codec_module_time_d4.py` uses `nn.InstanceNorm2d(..., affine=True)`, both on torch defaults. So this key is effectively ours; applying it uniformly is the coherent reading, and it is exactly faithful at the checkpoint's value. 3. The speech-enhancement manifest listed only `capabilities`, unlike the other audio workflow, which declares `adapter_abis`. The component already named the ABI, but a consumer reads the manifest to decide whether it can run the package at all. Declared — conditionally, since without STFT geometry no adapter ships and advertising one would state a requirement the caller does not have to meet. Both directions are tested. Not taken: the suggestion to read `compress_factor` from `stft_cfg` with a `model_cfg` fallback. The published nvidia/RE-USE config.json puts it in `model_cfg`, next to `hid_feature` — verified against the file on the Hub. Reading `stft_cfg` first is harmless but describes a layout upstream does not use, and inverting the order would silently fall back to the default for the real checkpoint. The field comment did say it belongs to the STFT front-end, which is what the finding picked up on; that is about where it is *applied*, not where it is *declared*, so the comment now says both and a test pins the source section. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 43bb31ac-c136-44bc-8cdf-70cbcedd4a42 Signed-off-by: Copilot <justinchuby@users.noreply.github.com>
I reported this model's throughput from a median of 3 runs. At 44.1 kHz the third run is the first steady one, so that median was still measuring warm-up and the figure I published was wrong. Re-measured over 9 runs, taking the median of runs 3 onward. The 8 kHz table barely moves and its steady window is tight (0.48-0.48 and 0.32-0.33 RTF), so those rows stand; they are updated to the re-measured values. What was missing is the two things that table cannot show: * First run costs ~4x steady (1066 ms and 2698 ms against 241 ms and 646 ms). That is one-time graph translation and kernel compilation, and it grows with the frequency extent — 7.0 s at 44.1 kHz. Anyone benchmarking this needs to iterate to the plateau, which is exactly the trap I fell into. * "Real time" was a claim about 8 kHz presented as a claim about the model. RE-USE is sampling-frequency-independent: it does not resample, it scales n_fft with the input rate, so 44.1 kHz means 883 frequency bins instead of 161 and ~5.5x the work per second of audio. Steady RTF is ~0.30 at 8 kHz, ~0.6 at 16 kHz, ~2.1 at 44.1 kHz. Real time up to 16 kHz, not beyond, on this hardware. Also replaced the "between 2x and 80x" range in TestExecutionProviderPartitioning with the specific consequence: an unclaimed Scan is a partition boundary, so a non-zero scan axis sends all 120 recurrences to CPU and gives up the whole 7.0x. The 80x half of that range came from an earlier ReverseSequence experiment whose rationale has since been rewritten in semantic terms, leaving the number unattributable. No code or test-logic changes; docstring and PR description only. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 43bb31ac-c136-44bc-8cdf-70cbcedd4a42 Signed-off-by: Copilot <justinchuby@users.noreply.github.com>
justinchuby
added a commit
to justinchuby/onnxruntime-mlx
that referenced
this pull request
Aug 23, 2026
## What The generic `Scan` handler unrolls its body once per timestep. For a Mamba-1 selective scan that is pathological: the RE-USE / SEMamba speech-enhancement model ([mobius#551](onnxruntime/mobius#551)) has **120 `Scan` nodes** with sequence lengths of 442 (frequency axis) and 256 (time axis), so the unroll emits on the order of half a million MLX nodes and the recurrence becomes **launch-bound** — ~14 tiny kernels per timestep, each moving only a few MB. This recognises that specific body and replaces the whole unroll with **one custom Metal kernel**. Each thread owns a single `(batch, d_inner)` pair and keeps its `d_state` accumulator in registers, so the running state is never materialised to device memory. Anything that does not match — a different body, an unsupported shape/dtype/`d_state`, a non-zero scan axis — falls through to the existing unroll unchanged. ## Why a kernel, and not ONNX primitives in mobius This was the main question in the investigation, and it has a concrete answer rather than a preference. Both alternatives were measured on the real checkpoint, not reasoned about. **1. The `cumsum` identity is unusable on this model.** The textbook form `h_t = exp(L_t) · cumsum(dBx_s · exp(−L_s))`, `L = cumsum(dt·A)`, is exact in fp64. But on the real weights `dt` reaches 2.78 and `A` reaches −46.4, so a **single** step's log-decay reaches **−129** — already past the fp32 `exp` range — and `exp(−L)` overflows to `inf`. Chunking does not rescue it; the worst in-chunk `|L|` is 1063 even at chunk 16: ``` dt : min 1.727e-07 max 2.781 A : min -46.380 max -0.230 per-step log decay min : -129.0 chunk C=16 -> worst in-chunk |L| = 1063 -> exp(-L) = inf (fp32) chunk C=64 -> worst in-chunk |L| = 4100 -> exp(-L) = inf (fp32) ``` **2. A division-free associative scan is numerically fine but slower than what we have.** Chunked Hillis-Steele over the `(a,b)` pairs has no division and no `exp` of a positive number, and it does agree to ~3e-07. But it has to materialise the `(chunk, batch, d_inner, d_state)` state, and `d_state=16` makes that a 16-fold bandwidth amplification. Measured on the real shapes (T=442, B=256, D=256, N=16), median per scan: | implementation | time | vs unrolled | |---|---|---| | **fused Metal kernel (this PR)** | **9.55 ms** | **7.8x faster** | | unrolled (status quo) | 74.26 ms | 1.0x | | associative, chunk 16 | 351.68 ms | 4.7x **slower** | | associative, chunk 64 | 470.18 ms | 6.3x **slower** | So the graph-layer route would have made the model *slower*. The recurrence is **bandwidth-bound, not depth-bound**, and trading sequential depth for memory traffic loses. The chunked-matmul trick real Mamba kernels use does not apply either: that is Mamba-2 / SSD, which works because `B` and `C` are shared within a head. In Mamba-1 `dA[d,n] = exp(dt[d]·a[d,n])` differs per channel, so the decay depends on `(t,k,d,n)` with no low-rank structure to contract `n` away early; it degenerates to O(T²·D·N). **The load-bearing point:** the fast form of this recurrence is *state resident in registers, never materialised*. That is precisely what ONNX primitives cannot express — every intermediate value in an ONNX graph is a tensor. This is why it belongs in the EP and not in the graph layer, despite the graph layer being the more portable place. ## Results M1 Max, 5.1 s file, median of **8 warm runs** (min–max in brackets). `first` is the cold run, which is what a single-shot CLI invocation actually pays. | rate | mode | first | median | range | |---|---|---|---|---| | 8 kHz | unrolled | 0.857 | 0.298 | 0.292–0.303 | | 8 kHz | **fused** | **0.201** | **0.160** | 0.159–0.166 | | 16 kHz | unrolled | 1.248 | 0.535 | 0.532–0.553 | | 16 kHz | **fused** | **0.355** | **0.310** | 0.307–0.323 | | 44.1 kHz | unrolled | 3.112 | 1.718 | 1.710–2.524 | | 44.1 kHz | **fused** | **0.948** | **0.862** | 0.849–0.894 | **44.1 kHz lands *around* realtime, and whether it clears the bar is load-dependent.** On the machine above the median is 0.862. An independent re-run by the reviewer, on the same hardware but with other work in flight, measured: | rate | mode | median | range | speedup | |---|---|---|---|---| | 8 kHz | unrolled → fused | 0.293 → **0.159** | 0.159–0.160 | 1.84x | | 44.1 kHz | unrolled → fused | 1.536 → **1.139** | 0.921–1.332 | 1.35x | 8 kHz reproduces almost exactly (1.84x vs 1.86x). 44.1 kHz does **not**: the fused median came out at 1.139, i.e. above realtime, with only the bottom of the range dipping under 1.0. Note the gap is on both sides — that run had a *faster* unrolled baseline (1.536 vs 1.718) and a *slower* fused result — so it is not simply a busier machine; the speedup itself varied (1.35x vs 1.99x). So: **8 kHz and 16 kHz clear realtime comfortably and reproducibly. 44.1 kHz sits at roughly 1.0 and should not be claimed as realtime.** Two independent measurements disagreeing across that line is exactly why it is stated this way rather than from the more favourable run. Two corrections to the original characterisation of the problem, both from measurement: * **Steady-state was never superlinear.** F′ goes 81 → 442 (5.46x) and steady RTF goes 0.298 → 1.718 (5.77x) — essentially linear. The superlinearity was in the **one-time** translate+compile cost (2.8 s → 7.1 s), which scales with the unrolled node count. Fusing collapses that too, which is why the cold run improves ~3.3x. * **Three runs is not enough to reach steady state at 44.1 kHz.** The plateau is only reached on the 3rd run; a median of 3 reports a partly-warm number. This is why the earlier figure looked worse than it is. Where the time actually goes, measured by replacing the `Scan` with a same-shape stand-in while keeping every projection (44.1 kHz): ``` full model steady : 8.50 s Scan bypassed : 2.96 s -> the 120 scans are 5.54 s = 65% ``` ## Correctness RE-USE output vs the **CPU EP on the real audio file**, with a wrap-aware angular difference for phase: | rate | | max abs delta amp | max wrapped abs delta phase | |---|---|---|---| | 8 kHz | fused | 9.775e-06 | 1.066e-03 | | 16 kHz | fused | 1.698e-05 | 3.051e-03 | | 44.1 kHz | fused | 2.560e-05 | 7.913e-03 | | 16 kHz | *unrolled (control)* | *1.635e-05* | *2.049e-03* | | 44.1 kHz | *unrolled (control)* | *2.185e-05* | *6.671e-03* | The control rows matter: the unrolled path **already on `main`** diverges from CPU by the same order. This is inherent MLX-vs-CPU fp32 divergence at near-silent bins accumulated through 30 blocks, not something the kernel introduces. Amplitude agreement stays in line with the 2.2e-05 baseline. The kernel computes `exp(dt·a)` per step exactly as the sequential reference does, so the −129 dynamic range is a non-issue: `exp` underflows to 0 and the state is annihilated, which is the benign direction and matches CPU. ## Tests `tests/ops/test_selective_scan.py`, and `tests/ops` stays green (**1349 passed / 64 skipped**, up from 1343/63). The important one is `test_scan_takes_the_fused_kernel_path`. The `Scan` is claimed by the EP either way and **the unrolled path returns the same answer**, so neither `assert_matches_cpu` nor a claim probe can tell the two apart — if mobius changes the emitted body and the pattern stops matching, every numeric test would stay green while performance silently reverts. That test therefore reads the EP trace and requires a `mlx_selective_scan` fast-path event with no composed `Scan`. A decline is recorded with a reason, so the fallback is never silent. Per the lesson from #46/#47, every new assertion was checked to actually fail: * forcing the unroll (`ONNXRUNTIME_EP_MLX_NO_SELECTIVE_SCAN=1`) → `test_scan_takes_the_fused_kernel_path` fails, reporting the decline reason. * sabotaging one term in the kernel → all 5 numeric tests fail; the path test still passes, confirming the two guard genuinely different things. Coverage includes the real checkpoint's dynamic range (`dt` to 2.78, `A` to −46.4), reverse `scan_input_directions`, a non-zero carried-in initial state, and fused-vs-unrolled agreement in-process. `ONNXRUNTIME_EP_MLX_NO_SELECTIVE_SCAN=1` forces the generic unroll, for A/B and bisection. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
…e-model Signed-off-by: Copilot <justinchuby@users.noreply.github.com> # Conflicts: # src/mobius/integrations/onnx_genai/__init__.py
Member
Author
|
Needs a metadata review |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Adds NVIDIA's RE-USE speech-enhancement generator (SEMamba): a dense conv encoder, 30 time-frequency bidirectional Mamba blocks, and separate magnitude/phase decoders (9.61M params).
It isn't a transformers architecture — no
model_type, and the checkpoint targetsmamba_ssm— so neitherAutoConfignor the existing decode-time Mamba components applied.What's new
SequenceSelectiveScan— the Mamba-1 recurrence over a whole sequence as a single ONNXScan.dA/dBxare derived inside the scan body rather than materialised for all timesteps, which keeps the working set at(batch, seq, d_inner)instead of(batch, seq, d_inner, d_state)(~500 MB per Mamba module at 2 s of audio).SequenceMambaBlock— stateless full-sequence counterpart toMambaBlock, sharing its parameter names.Conv2dgains an optional trailingdilationargument.models/reuse.py—ReUseConfig, the encoder/decoder stack, and a hand-rolledatan2(ONNX has noAtan2).SpeechEnhancementTask—noisy_mag/noisy_phain,denoised_mag/pha/comout. STFT/ISTFT stay outside the graph.build_reuse()— loads config + weights from a directory or the Hub, sincebuild()goes throughAutoConfigand can't discover this repo.Module attribute names are structured so encoder/decoder parameters land on the checkpoint's
nn.Sequentialindices directly;preprocess_weightsonly nests the flat SSM parameters underssm.Correctness
Against a pure-PyTorch transcription of the reference model with the real weights: magnitude and complex outputs match to 4.2e-06, all 1429 initializers populated.
Phase exceeds 1e-04 at 10 of 6601 points — those are exactly the smallest-radius points, where
atan2is ill-conditioned near the origin. Error in the complex plane there is ≤1.7e-05, so it's float32 conditioning, not a logic error.Performance
The
Scanis fed time-major so plugin EPs that require a scan axis of 0 can claim it. Measured againstonnxruntime-ep-mlx0.29.4 (M1 Max, 9 runs, median of runs 3+ — see the warm-up note below):At 2 s the model goes from well under real time to comfortably over it. The scan length is only static when the graph inputs are, so this benefits fixed-size/chunked export; RE-USE ships
inference_chunk.py.Two caveats worth stating, because neither is visible in that table:
Steady state only; the first run costs much more. MLX first-run times for the rows above are 1066 ms and 2698 ms — 4x the steady figure. The gap is one-time graph translation and kernel compilation, and it grows with the frequency extent (7.0 s at 44.1 kHz). Benchmark this model with enough iterations to reach the plateau: at 44.1 kHz the third run is the first steady one, so a median-of-3 silently reports a warm-up number.
Real time is a statement about 8 kHz, not about the model. RE-USE is sampling-frequency-independent: it does not resample, it scales
n_fftwith the input rate, so 44.1 kHz means 883 frequency bins instead of 161 and about 5.5x the work per second of audio. Measured steady RTF is ≈0.30 at 8 kHz, ≈0.6 at 16 kHz, and ≈2.1 at 44.1 kHz — real time up to 16 kHz, not beyond it on this hardware.These figures predate the EP-side scan fusion. The version pinned above unrolls each
Scaninto one graph copy per timestep. justinchuby/onnxruntime-mlx#49 has since fused the Mamba-1 selective scan into a single Metal kernel, which roughly halves these numbers (8 kHz RTF 0.29 → 0.16 on the same machine) and cuts the first-run cost by ~4x. Nothing in this PR depends on it — the graph is unchanged and the fused path is recognised from the existingScan— but quote the version when citing throughput, since the two differ by about 2x.Testing
45 new tests: graph construction, config extraction, weight-name alignment, ORT execution across input lengths,
atan2over all quadrants, and onnx-genai metadata validated against onnx-genai's committed JSON schema (with and without the preprocessing program). RE-USE is also driven bySPEECH_CONFIGSfor L1/L3 coverage.Full suite: no new failures against the branch point.
lintrunnerclean.Side effects
Exporting this surfaced two upstream shape-inference bugs, both fixed: justinchuby/onnx-shape-inference#119 (anonymous
SymbolicDimintoScan/Loopbodies) and justinchuby/onnx-shape-inference#120 (Max/Minsymbolic value retention). With justinchuby/onnx-shape-inference#119 applied, inference succeeds here and the outputs carry real symbolic expressions.Follow-up
The encoder's output frequency extent is statically derivable from
n_fftbut is currently computed viaShape. Making it a build-time constant would let an EP claim the frequency-axis scans even with a fully dynamic time axis.