Fix JAX convolution dimension numbers

Author: prishajain1

src/maxdiffusion/models/ltx2/vocoder_ltx2.py

@@ -100,7 +100,7 @@ def __call__(self, x: Array) -> Array:
         filter_expanded,
         window_strides=(self.ratio,),
         padding="VALID",
-        dimension_numbers=('NLC', 'WIO', 'NLC'),
+        dimension_numbers=('NLC', 'LIO', 'NLC'),
         feature_group_count=num_channels,
     )
     return x_filtered
@@ -156,7 +156,7 @@ def __call__(self, x: Array) -> Array:
         window_strides=(1,),
         padding="VALID",
         lhs_dilation=(self.ratio,),
-        dimension_numbers=('NLC', 'WIO', 'NLC'),
+        dimension_numbers=('NLC', 'LIO', 'NLC'),
         feature_group_count=num_channels,
     )
 
@@ -486,7 +486,7 @@ def __call__(self, waveform: Array) -> tuple[Array, Array]:
         self.forward_basis.value,
         window_strides=(self.hop_length,),
         padding="VALID",
-        dimension_numbers=('NLC', 'WIO', 'NLC'),
+        dimension_numbers=('NLC', 'LIO', 'NLC'),
     )
 
     n_freqs = spec.shape[-1] // 2