Fix vectorization for stft

[jonatanklosko]

Currently vectorization results in an error:

** (CompileError) /Users/jonatanklosko/git/nx_signal/lib/nx_signal.ex:359: the do-block in while must return tensors with the same shape, type, and names as the initial arguments.

{
 <<<<< Body (do-block) <<<<<
 #Nx.Tensor<
   vectorized[batch: 1]
   f32[3001][400]
 >
 ==========
 #Nx.Tensor<
   f32[3001][400]
 >
 >>>>>     Initial     >>>>>
 , #Nx.Tensor<
   s64
 >, #Nx.Tensor<
   s64
 >, #Nx.Tensor<
   vectorized[batch: 1]
   f32[480400]
 >}

Is this something we should relax in Nx? cc @josevalim

[josevalim]

I'd prefer to keep them the same as vectorized and not will lead to different semantics. @polvalente decides though.

[josevalim]

Btw, is this something that should have been caught by a test?

[polvalente]

I think that Nx is correct and this fix is needed here in NxSignal as-is :)

[polvalente]

Btw, is this something that should have been caught by a test?

Maybe, but we don't generally test for vectorization on stuff outside of Nx.

For instance, there's no explicit support for vectorization in Scholar nor Axon

[josevalim]

i mean, not really testing the vectorization, but to test stft works with vectorized tensors which would help regressions on this code in the future.