Repository: pytorch
Issue: torch.export torchaudio kaldi module
Labels: triaged,enhancement,oncall: pt2,export-triaged,oncall: export

### 🚀 The feature, motivation and pitch

# Preface

Hello! I am trying to export the codes for audio preprocessing on `kaldi` module in torchaudio.

I have few feature requests:
- Is there any plans to support `torch.map` to support CUDA tensors? 
- Is there a plan to support `ta_kaldi.fbank` to be batch-compatible, or `torch.export`'able in general?


### Alternatives

_No response_

### Additional context

I was trying to use `torch.export` to extract the AOT graph of the computation, as following:

```python
import torch
import torchaudio.compliance.kaldi as ta_kaldi

def preprocess_forloop(
    source: torch.Tensor,
    fbank_mean: float = 15.41663, #-12
    fbank_std: float = 6.55582, #4
) -> torch.Tensor:

    if len(source.shape)!=2: # get sequence data, B x L
        return source.transpose(1,2)

    fbanks = []
    for waveform in source:
        waveform = waveform.unsqueeze(0)  * 2 ** 15
        fbank = ta_kaldi.fbank(waveform, num_mel_bins=128, sample_frequency=48000, frame_length=25, frame_shift=10)
        fbanks.append(fbank)

    fbank = torch.stack(fbanks, dim=0)

    fbank = (fbank - fbank_mean) / (2 * fbank_std)
    return fbank

class AudioPreprocessModel(torch.nn.Module):

    def forward(self, source: torch.Tensor) -> torch.Tensor:
        # source: [B, 720000]
        # return: [B, 1498, 128]
        return preprocess_forloop(source)

inp = torch.zeros(2, 720000)
bs = torch.export.Dim("bs")
ep = torch.export.export(AudioPreprocessModel(), (inp,), dynamic_shapes={"source": {0: bs}})

```

This code snippet is from BEATs, a popular audio foundation model. The preprocess_forloop function is brought from [preprocess](https://github.com/microsoft/unilm/blob/master/beats/BEATs.py#L118-L131) in BEATs repo.

The code uses `torchaudio.compliance.kaldi.fbank` to extract Mel spectrogram from the raw audio signal

The `fbank` function accepts `[n_channels, T]` tensor and generates `[T_2, mel_bins]` tensor.

# Issues

I tried to `torch.export` this preprocessing function, but encountered multiple obstacles.

🟡: able to get around with patching

🔴: weren’t able to solve

## 1. 🔴 for loop inside the function

The very nature of `fbank` is that it is not batch-compatible, meaning we have to run for loop within a batch to extract spectrograms.

The presence of for loop makes `torch.export` to be challenging, since TorchDynamo graph capturing attempts to unroll the loop, which can vary depending on the batch size. The graph capturing and subsequent `torch.export` fails with this code.

We can use [torch.map](https://pytorch.org/docs/stable/generated/exportdb/torch.map.html) to alleviate this issue, something like as follows.

```python
from functorch.experimental.control_flow import map as torch_map
def preprocess_torchmap(
    source: torch.Tensor,
    fbank_mean: float = 15.41663, #-12
    fbank_std: float = 6.55582, #4
) -> torch.Tensor:
    def fbank_body(waveform: torch.Tensor) -> torch.Tensor:
        waveform = waveform.unsqueeze(0)  * 2 ** 15
        # [1, 720000] -> [1498, 128]
        fb = ta_kaldi.fbank(waveform, num_mel_bins=128, sample_frequency=48000, frame_length=25, frame_shift=10)
        return fb
    fbank = torch_map(fbank_body, source)
    fbank = (fbank - fbank_mean) / (2 * fbank_std)
    return fbank

```

However the `torch.map` didn’t supprt CUDA tensors, so I eventually weren’t able to export the function.

## 2. 🟡 Intermediate tensor generation

The [`_get_epsilon`](https://github.com/pytorch/audio/blob/main/src/torchaudio/compliance/kaldi.py#L35-L36) function encountered export failure, as it attempted to make a intermediate tensor with unknown value ahead of time.

Resolved with `torch._dynamo.assume_constant_result` decorator, with following monkey patching:

```python
ta_kaldi._get_epsilon = torch._dynamo.assume_constant_result(ta_kaldi._get_epsilon)

```

## 3. 🟡 Symbolic int does not support python native functions

The [`_next_power_of_2`](https://github.com/pytorch/audio/blob/main/src/torchaudio/compliance/kaldi.py#L39-L41) function encountered export failure, because the symbolic representation of python `int` didn’t support `.bit_length()` operation.

Resolved with `math.log2` function, albeit being much inefficient computation.

```python
def _next_power_of_2(x: int) -> int:
    r"""Returns the smallest power of 2 that is greater than x"""
    return 1 if x == 0 else 2 ** math.ceil(math.log2(x))
ta_kaldi._next_power_of_2 = _next_power_of_2

```

## 4. 🟡 `torch.tensor` generation within torch map

The creation of tensor with `torch.tensor` inside the `torch.map` body creates a failure mode at a graph verification phase.

Attached a repro code [gist](https://gist.github.com/swmoon00/536038b1291a1b1970614757a903edf0) snippet.

I substituted `torch.tensor` creation with `torch.ones` creation, in the [`_get_log_energy`](https://github.com/pytorch/audio/blob/main/src/torchaudio/compliance/kaldi.py#L116-L122) function which caused failing.

```python
def _get_log_energy(strided_input: torch.Tensor, epsilon: torch.Tensor, energy_floor: float) -> torch.Tensor:
    r"""Returns the log energy of size (m) for a strided_input (m,*)"""
    device, dtype = strided_input.device, strided_input.dtype
    log_energy = torch.max(strided_input.pow(2).sum(1), epsilon).log()  # size (m)
    if energy_floor == 0.0:
        return log_energy

    log_energy_flow_tensor = torch.ones(1, device=device, dtype=dtype)[0] * math.log(energy_floor)
    return torch.max(log_energy, log_energy_flow_tensor)
    # return torch.max(log_energy, torch.tensor(math.log(energy_floor), device=device, dtype=dtype))
ta_kaldi._get_log_energy = _get_log_energy

```

## 5. 🔴 `torch.fft.rfft` does not support `bfloat16`

This is related to complex number support on torch library function. Not straightforward to fix it.
Created a separated [issue](https://github.com/pytorch/pytorch/issues/139313) on this.

cc @ezyang @chauhang @penguinwu @avikchaudhuri @gmagogsfm @zhxchen17 @tugsbayasgalan @angelayi @suo @ydwu4