A limited port of scipy.signal.spectral (from SciPy 1.5.4) to JAX.
Supports:
Note: STFT only works when nfft == nperseg and nperseg (window size) is a multiple of nperseg - noverlap (hop size).
See also the official jax.scipy.signal (currently missing an implementation of spectral).
pip install 'git+https://github.com/cifkao/jax-spectral.git@main#egg=jax-spectral'