[feature request] support exporting the model to ONNX format

[opringle]

Currently attempting to export a model that uses KeyedJaggedTensor to ONNX format results in the following error:

gm = symbolic_trace(model)
scripted_gm: torch.jit.ScriptModule | torch.jit.ScriptFunction = torch.jit.script(gm)  # type: ignore

torch.onnx.export(
    model=scripted_gm,
    args=...,
    f="model.onnx",
    export_params=True,
    verbose=True,
    opset_version=10,
    do_constant_folding=True,
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={
        "input": {0: "batch_size"},
        "output": {0: "batch_size"},
    },
)

RuntimeError: Only tuples, lists and Variables are supported as JIT inputs/outputs. Dictionaries and strings are also accepted, but their usage is not recommended. Here, received an input of unsupported type: KeyedJaggedTensor

[henrylhtsang]

Linking related issue: https://github.com/pytorch/torchrec/issues/1221

[opringle]

Is this a feature request we should make to the onnx team or torchrec team?