mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-17 20:36:05 +00:00
Code Issues Packages Projects Releases Wiki Activity

Remove use of int casting in STFT collapse of batch dimensions.

Browse Source 
PiperOrigin-RevId: 524115535
This commit is contained in:
jax authors 2023-04-13 15:14:37 -07:00
parent 3e93833ed8
commit 0fd5b2ca61
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
jax/_src/scipy/signal.py
View File

@ -233,7 +233,7 @@ def _fft_helper(x: Array, win: Array, detrend_func: Callable[[Array], Array],
else:
step = nperseg - noverlap
batch_shape = list(batch_shape)
x = x.reshape((math.prod(batch_shape)), signal_length)[..., np.newaxis]
x = x.reshape((math.prod(batch_shape), signal_length, 1))
result = jax.lax.conv_general_dilated_patches(
x, (nperseg,), (step,),
'VALID',