(
op: torch._ops.OpOverload,
fn: Fn,
args: Argument,
kwargs: Dict[str, Argument],
)
| 811 | |
| 812 | |
| 813 | def map_args( |
| 814 | op: torch._ops.OpOverload, |
| 815 | fn: Fn, |
| 816 | args: Argument, |
| 817 | kwargs: Dict[str, Argument], |
| 818 | ) -> Tuple[Argument, Dict[str, Argument]]: |
| 819 | assert isinstance(args, tuple) |
| 820 | assert isinstance(kwargs, dict) |
| 821 | args = list(args) |
| 822 | kwargs = kwargs.copy() |
| 823 | |
| 824 | def update(key: K, args: MutableMapping[K, PyTree], schema: ArgSchema) -> None: |
| 825 | args[key] = fn(args[key], schema) |
| 826 | |
| 827 | for i, schema in enumerate(op._schema.arguments): |
| 828 | if schema.name in kwargs: |
| 829 | update(schema.name, kwargs, schema) |
| 830 | elif not schema.kwarg_only and i < len(args): |
| 831 | update(i, args, schema) # pyre-ignore |
| 832 | |
| 833 | return tuple(args), kwargs |
no test coverage detected