More intelligent and flexible input and output tensor naming for TorchScript Exporter

[EthanMarx]

Previously, the TorchScript exporter directly enforced input and output tensor names follow the INPUT__{i} OUTPUT__{i} convention. The main reason for this is that triton expects certain naming conventions to deal with the fact that TorchScript models don't come with some necessary metadata triton needs to match input tensors to the appropriate ordering expected by the model. If your model only has one input and one output, this detail is irrelevant, and any naming will work.

This convention is not the only way to deal with this (see here): one can also name the tensors by the parameter names in the models forward, or name the tensors with <name>__<index>.

In any case, the INPUT__{i} OUTPUT__{i} tensor naming enforcement was leading to some difficulties when ensembling multiple torchscript models due to conflicting tensor names.

This PR makes naming tensors in TorchScript models more flexible by doing the following:

• Input tensors

  1. If the user passes a dictionary of {tensor_name: shape} for the input_shapes, the user passed names will be used. A warning will be given telling users to be careful when using their own name.
  2. If the user passes a List[shape] to input_shapes, the names will be inferred directly from the forward method of the TorchScript model - this behavior is recommended by triton.

• Output tensors

  1. If the user passes a List[name] to output_names, the user passed names will be used
  2. If output_names is left as None, the OUTPUT__{i} convention will be used