Modify some operations to be compatible for torch.onnx.export

[robin-p-schmitt]

I am currently trying to export a CTC model based on the RF Conformer with torch.onnx.export. For this, some operations need to be modified:

• ONNX does not support torch.full taking a boolean value. Instead, I cast the boolean value to float and convert the resulting tensor back to bool.
• ONNX does not support torch.stft returning a complex dtype. Instead, I set return_complex=False and handle the resulting tensor accordingly.
• The most extensive issue is that ONNX does not support torch.reshape to take a shape with tensor elements. In the torch backend, I changed all torch.reshape calls to TorchBackend.reshape_raw and do different checks to handle the logic without torch.reshape. So far, I only included the cases which I encountered during my specific ONNX export. I think the checks could maybe be designed in a smarter way which is why I marked this as a draft for now.

Right now, I am stuck with one of the cases in TorchBackend.reshape_raw: elif len(raw_tensor.shape) - len(shape) == -1:. This is true, for example, when a single dimension is split into two dimensions which are given as tensors. I am not sure whether this case can be generally handled here. Any ideas?

[Gerstenberger]

Just some thoughts in general.
In my view, ideally, onnx export problems should be fixed in PyTorch itself.

FYI, torch.onnx.export also supports adding new or overwriting existing torch op to onnx conversions via custom_translation_table parameter and OnnxScript. This would be more robust, because you would not have to search for instances in returnn. So this could also be done with a custom export step for the torch conversion script.