Rewrite the CHIRP model code so it can be lowered to a supported Stab…

…leHLO/MHLO IR. PiperOrigin-RevId: 723162070
google-research · Feb 4, 2025 · 1a49cfb · 1a49cfb
1 parent 827f0e1
commit 1a49cfb
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 8 deletions.
diff --git a/chirp/audio_utils.py b/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(

diff --git a/chirp/train_tests/frontend_test.py b/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": {