Skip to content

core/ops/fft: Stft::axes_mapping unlocks pulsification on non-STFT axes#2331

Merged
JulienBalianSonos merged 1 commit into
mainfrom
feat/stft-pulse-non-stft-axes
Jun 2, 2026
Merged

core/ops/fft: Stft::axes_mapping unlocks pulsification on non-STFT axes#2331
JulienBalianSonos merged 1 commit into
mainfrom
feat/stft-pulse-non-stft-axes

Conversation

@JulienBalianSonos
Copy link
Copy Markdown
Collaborator

@JulienBalianSonos JulienBalianSonos commented Jun 2, 2026

Summary

Stft is the only typed op left in core/src/ops/fft.rs without an axes_mapping. Add one. The mapping is custom because Stft is not rank-preserving: it inserts a frame axis at op.axis + 1.

The mapping declares:

  • leading dims (0..op.axis): 1-to-1 input ↔ output (natural).
  • input op.axis (time) ↔ output op.axis (n_frames, same position, smaller dim).
  • output op.axis + 1: output-only (the inserted frame axis).
  • trailing input dims (op.axis + 1..rank, including the complex pair) ↔ output dims shifted right by 1.

Side effect: PulseWrappingOp can now handle a non-STFT streaming axis without bailing on could not track pulsing axis. I don't have an in-the-wild model that needs this today; the typical STFT use case streams the time axis, which stft_pulsify already handles. The patch is for structural completeness (every TypedOp should declare its axes contract) and symmetry with the Fft fix in #2261.

@JulienBalianSonos JulienBalianSonos merged commit 8398ec0 into main Jun 2, 2026
55 checks passed
@JulienBalianSonos JulienBalianSonos deleted the feat/stft-pulse-non-stft-axes branch June 2, 2026 15:01
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants