torch.bitwise_right_shift和bitwise_and导出到onnx算子的问题

代码中使用了bitpack来量化-反量化，其中用到了torch.bitwise_right_shift和bitwise_and函数，在导出ONNX模型时报错，特此记录。

1. Opset version选择

从ONNX DOCUMENT可以看到，onnx直到opset 18才添加了对BitwiseAnd的支持。

2. 直接导出，报错

导出脚本：

torch.onnx.export(model, (x), "test.onnx", opset_version=18, verbose=True, input_names=['input'])

运行后：

    raise errors.UnsupportedOperatorError(
torch.onnx.errors.UnsupportedOperatorError: Exporting the operator 'aten::bitwise_right_shift' to ONNX opset version 18 is not supported. Please feel free to request support or submit a pull request on PyTorch GitHub: https://github.com/pytorch/pytorch/issues.

可以看到虽然ONNX 官方说了支持，但并没有与torch函数对应起来。

3. 添加函数支持

找到torch安装位置，打开torch/onnx/symbolic_opset18.py，添加下列代码：

@_onnx_symbolic("aten::bitwise_right_shift")
@_beartype.beartype
def bitwise_right_shift(g, input: torch.Value, other: torch.Value):
    return g.op("BitShift", input, other, direction_s='RIGHT')

@_onnx_symbolic("aten::bitwise_and")
@_beartype.beartype
def bitwise_and(g, input: torch.Value, other: torch.Value):
    return g.op("BitwiseAnd", input, other)

4. 运行

可能会遇到一些数据类型的报错，根据ONNX文档修改torch运行过程中涉及到tensor的dtype即可.