Method replace

tensorrt_llm/graph_rewriting.py:381–390

(arg: Any)

        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)