Repository: pytorch
Issue: EMFORMER_RNNT not compilable
Labels: feature,triaged,oncall: pt2,module: dynamic shapes

### 🚀 The feature

## Repro

```python
import torch
import torchaudio
from torchaudio.io import StreamReader

bundle = torchaudio.pipelines.EMFORMER_RNNT_BASE_LIBRISPEECH

feature_extractor = bundle.get_streaming_feature_extractor()
decoder = bundle.get_decoder()
token_processor = bundle.get_token_processor()

# random values
# This works
decoder(torch.randn(80, 80), length=torch.tensor([1.0]), beam_width=10)

decoder = torch.compile(decoder, fullgraph=True)

# This does not work
decoder(torch.randn(80, 80), length=torch.tensor([1.0]), beam_width=10)
```

## Error

```
(nightly) marksaroufim@marksaroufim-mbp zzz % ls
add.ff                          model.ff                        test.py                         test4.py
buck2-aarch64-apple-darwin.zst  tes2.py                         test3.py
(nightly) marksaroufim@marksaroufim-mbp zzz % 
(nightly) marksaroufim@marksaroufim-mbp zzz % python test4.py
Intel MKL WARNING: Support of Intel(R) Streaming SIMD Extensions 4.2 (Intel(R) SSE4.2) enabled only processors has been deprecated. Intel oneAPI Math Kernel Library 2025.0 will require Intel(R) Advanced Vector Extensions (Intel(R) AVX) instructions.
Intel MKL WARNING: Support of Intel(R) Streaming SIMD Extensions 4.2 (Intel(R) SSE4.2) enabled only processors has been deprecated. Intel oneAPI Math Kernel Library 2025.0 will require Intel(R) Advanced Vector Extensions (Intel(R) AVX) instructions.
/Users/marksaroufim/opt/anaconda3/envs/nightly/lib/python3.10/site-packages/executorch/extension/pytree/__init__.py:25: UserWarning: Unable to import executorch.extension.pytree, using native torch pytree instead.
  warnings.warn(
[2023-07-26 18:57:18,049] torch._dynamo.convert_frame: [WARNING] WON'T CONVERT forward /Users/marksaroufim/opt/anaconda3/envs/nightly/lib/python3.10/site-packages/torchaudio/models/rnnt_decoder.py line 267 
[2023-07-26 18:57:18,049] torch._dynamo.convert_frame: [WARNING] due to: 
[2023-07-26 18:57:18,049] torch._dynamo.convert_frame: [WARNING] Traceback (most recent call last):
[2023-07-26 18:57:18,049] torch._dynamo.convert_frame: [WARNING]   File "/Users/marksaroufim/opt/anaconda3/envs/nightly/lib/python3.10/site-packages/torch/_inductor/fx_passes/split_cat.py", line 125, in normalize_cat_default
[2023-07-26 18:57:18,049] torch._dynamo.convert_frame: [WARNING]     assert all(ndim == x.meta["example_value"].dim() for x in tensors)
[2023-07-26 18:57:18,049] torch._dynamo.convert_frame: [WARNING] torch._dynamo.exc.BackendCompilerFailed: backend='inductor' raised:
[2023-07-26 18:57:18,049] torch._dynamo.convert_frame: [WARNING] AssertionError: 
[2023-07-26 18:57:18,049] torch._dynamo.convert_frame: [WARNING] 
[2023-07-26 18:57:18,049] torch._dynamo.convert_frame: [WARNING] 
[2023-07-26 18:57:18,049] torch._dynamo.convert_frame: [WARNING] converting frame raised error, suppressing error
[2023-07-26 18:57:21,172] torch._dynamo.convert_frame: [WARNING] WON'T CONVERT transcribe /Users/marksaroufim/opt/anaconda3/envs/nightly/lib/python3.10/site-packages/torchaudio/models/rnnt.py line 580 
[2023-07-26 18:57:21,172] torch._dynamo.convert_frame: [WARNING] due to: 
[2023-07-26 18:57:21,172] torch._dynamo.convert_frame: [WARNING] Traceback (most recent call last):
[2023-07-26 18:57:21,172] torch._dynamo.convert_frame: [WARNING]   File "/Users/marksaroufim/opt/anaconda3/envs/nightly/lib/python3.10/site-packages/torch/_inductor/fx_passes/split_cat.py", line 125, in normalize_cat_default
[2023-07-26 18:57:21,172] torch._dynamo.convert_frame: [WARNING]     assert all(ndim == x.meta["example_value"].dim() for x in tensors)
[2023-07-26 18:57:21,172] torch._dynamo.convert_frame: [WARNING] torch._dynamo.exc.BackendCompilerFailed: backend='inductor' raised:
[2023-07-26 18:57:21,172] torch._dynamo.convert_frame: [WARNING] AssertionError: 
[2023-07-26 18:57:21,172] torch._dynamo.convert_frame: [WARNING] 
[2023-07-26 18:57:21,172] torch._dynamo.convert_frame: [WARNING] 
[2023-07-26 18:57:21,172] torch._dynamo.convert_frame: [WARNING] converting frame raised error, suppressing error
[2023-07-26 18:57:24,048] torch._dynamo.convert_frame: [WARNING] WON'T CONVERT forward /Users/marksaroufim/opt/anaconda3/envs/nightly/lib/python3.10/site-packages/torchaudio/models/rnnt.py line 220 
[2023-07-26 18:57:24,048] torch._dynamo.convert_frame: [WARNING] due to: 
[2023-07-26 18:57:24,048] torch._dynamo.convert_frame: [WARNING] Traceback (most recent call last):
[2023-07-26 18:57:24,048] torch._dynamo.convert_frame: [WARNING]   File "/Users/marksaroufim/opt/anaconda3/envs/nightly/lib/python3.10/site-packages/torch/_inductor/fx_passes/split_cat.py", line 125, in normalize_cat_default
[2023-07-26 18:57:24,048] torch._dynamo.convert_frame: [WARNING]     assert all(ndim == x.meta["example_value"].dim() for x in tensors)
[2023-07-26 18:57:24,048] torch._dynamo.convert_frame: [WARNING] torch._dynamo.exc.BackendCompilerFailed: backend='inductor' raised:
[2023-07-26 18:57:24,048] torch._dynamo.convert_frame: [WARNING] AssertionError: 
[2023-07-26 18:57:24,048] torch._dynamo.convert_frame: [WARNING] 
[2023-07-26 18:57:24,048] torch._dynamo.convert_frame: [WARNING] 
[2023-07-26 18:57:24,048] torch._dynamo.convert_frame: [WARNING] converting frame raised error, suppressing error
```

### Motivation, pitch

It's a popular model folks have been asking me about and I'd like get it deployed to a mobile device, starting with the decoder

Here's my progress so far
* torch.compile was graph breaking on jit annotate


```
@@ -236,7 +236,7 @@ class RNNTBeamSearch(torch.nn.Module):
         b_hypos = self._init_b_hypos(device) if hypo is None else hypo
         for t in range(n_time_steps):
             a_hypos = b_hypos
-            b_hypos = torch.jit.annotate(List[Hypothesis], [])
+            b_hypos = []
             key_to_b_hypo: Dict[str, Hypothesis] = {}
             symbols_current_t = 0
 
@@ -292,7 +292,7 @@ class RNNTBeamSearch(torch.nn.Module):
         enc_out, _ = self.model.transcribe(input, length)
         return self._search(enc_out, None, beam_width)
 
-    @torch.jit.export
+    # @torch.jit.export
     def infer(
         self,
         input: torch.Tensor,
```

### Alternatives

_No response_

### Additional context

_No response_

cc @ezyang @bdhirsh @anijain2305 @zou3519 @chauhang @wconstab @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @ipiszy @chenyang78 @aakhundov