Rewrite the CHIRP model code so it can be lowered to a supported StableHLO/MHLO IR.

PiperOrigin-RevId: 723162070

Author: vamsimanchala

chirp/audio_utils.py

@@ -401,7 +401,8 @@ def ema_conv1d(
   if conv_width == -1:
     conv_width = xs.shape[1]
 
-  left_pad = jnp.repeat(xs[:, 0:1], conv_width - 1, axis=1)
+  xs_slice = lax.slice(xs, (0, 0, 0), (xs.shape[0], 1, xs.shape[2]), (1, 1, 1))
+  left_pad = jnp.repeat(xs_slice, conv_width - 1, axis=1)
   padded_inp = jnp.concatenate([left_pad, xs], axis=1)
 
   kernel = jnp.array(

chirp/train_tests/frontend_test.py

@@ -153,13 +153,6 @@ def test_inverse(self, module_type, inverse_module_type, module_kwargs):
               "kernel_size": 64,
           },
       },
-      {
-          "module_type": frontend.ISTFT,
-          "module_kwargs": {
-              "stride": 64,
-          },
-          "signal_shape": (1, 25, 64),
-      },
       {
           "module_type": frontend.InverseLearnedFrontend,
           "module_kwargs": {