Method replace_input_with

tensorrt_llm/graph_rewriting.py:373–392

Replace the input `src` with the input `dst` in the raw_inputs. src: Tensor dst: Tensor

(self, src, dst)

        return copy(self.raw_inputs)

    def replace_input_with(self, src, dst):
        """Replace the input `src` with the input `dst` in the raw_inputs.

        src: Tensor
        dst: Tensor
        """
        from .functional import Tensor

        def replace(arg: Any):
            if isinstance(arg, Tensor):
                if arg.trt_tensor is src.trt_tensor:
                    return dst
                return arg
            elif isinstance(arg, (list, tuple)):
                return [replace(x) for x in arg]
            elif isinstance(arg, dict):
                return {k: replace(v) for k, v in arg.items()}
            return arg

        replace(self.raw_inputs)

    def replace_outputs_uses_with(self, net: Network, new_outs: List[Any]):
        """Replace the output users with the new outputs.