A helper function to update an argument container without changing it. This can be used with both args and kwargs.
(
args: Tuple[Argument, ...], key: int, val: torch.fx.Node
)
| 105 | |
| 106 | |
| 107 | def update_args( |
| 108 | args: Tuple[Argument, ...], key: int, val: torch.fx.Node |
| 109 | ) -> Tuple[Argument, ...]: |
| 110 | """ |
| 111 | A helper function to update an argument container without changing it. |
| 112 | This can be used with both args and kwargs. |
| 113 | """ |
| 114 | if isinstance(args, dict): |
| 115 | new_dict = copy.copy(args) |
| 116 | new_dict[key] = val |
| 117 | return new_dict |
| 118 | |
| 119 | assert isinstance(args, tuple) |
| 120 | new_tuple = list(args) |
| 121 | new_tuple[key] = val |
| 122 | return tuple(new_tuple) |
| 123 | |
| 124 | |
| 125 | class DebugPass(PassBase): |