Description
Add support for operator.get
and torch.ops.aten._to_copy
and torch.ops.aten.clone
For operator.get
we need an evaluator-style implementation which simply dispatches the call out to the Python implementation. For torch.ops.aten._to_copy
, the implementation will be a bit more challenging, since we can only support a subset of these operations. Currently, a simple way to perform the differentiation of which _to_copy
ops we can handle is to modify the support logic to check if the to
converter has a cast which can be performed in TRT. An example of an allowed cast might be torch.float -> torch.int
, whereas torch.int ->torch.long
is not allowed:
TensorRT/py/torch_tensorrt/dynamo/backend/lowering/_partition.py
Lines 109 to 131 in 6e4aa0b
Long Term Solution
A long term solution to this issue of selective conversion is to alter the CONVERSION
dictionary to also store an auxiliary function in addition to its converter. This function will take an input node of the converter's type and return whether that node can be converted. For most converters, this function will be: lambda node: True
. For _to_copy
, however, it might be something like:
def converts_to_copy(node: torch.fx.Node):
return "dtype" in kwargs and kwargs["dtype"] in APPROVED_CAST_DTYPES