(lowered_module, *args)
| 77 | @executorch_call_delegate.py_impl(torch._C.DispatchKey.CompositeExplicitAutograd) |
| 78 | # pyre-ignore |
| 79 | def call_delegate_cpu(lowered_module, *args): |
| 80 | # FX creates this immutable_dict/list concept. Get rid of this. |
| 81 | map_types = { |
| 82 | torch.fx.immutable_collections.immutable_dict: dict, |
| 83 | torch.fx.immutable_collections.immutable_list: list, |
| 84 | } |
| 85 | new_args = pytree.tree_map_only( |
| 86 | tuple(map_types.keys()), |
| 87 | lambda a: map_types[type(a)](a), |
| 88 | args, |
| 89 | lambda a: isinstance(a, tuple(map_types.keys())), |
| 90 | ) |
| 91 | return lowered_module.original_module.module()(*new_args) |
| 92 | |
| 93 | @executorch_call_delegate.py_impl(torch._C.DispatchKey.Autograd) |
| 94 | # pyre-ignore |
no test coverage detected