Searched refs:module_call_specs (Results 1 – 2 of 2) sorted by relevance
| /aosp_15_r20/external/pytorch/torch/_export/ |
| H A D | wrappers.py | 61 def _wrap_submodule(mod, path, module_call_specs): argument
71 if path in module_call_specs:
72 assert module_call_specs[path]["in_spec"] == in_spec
73 assert module_call_specs[path]["out_spec"] == out_spec
74 module_call_specs[path] = {"in_spec": in_spec, "out_spec": out_spec}
|
| /aosp_15_r20/external/pytorch/torch/export/ |
| H A D | _trace.py | 132 module_call_specs: Dict[str, Dict[str, pytree.TreeSpec]]
556 module_call_specs: Dict[str, Dict[str, pytree.TreeSpec]] = {}
558 f, preserve_module_call_signature, module_call_specs
584 gm_torch_level.meta["module_call_specs"] = module_call_specs
1074 module_call_specs: Dict[
1076 ] = export_artifact.module_call_specs
1081 for fqn, specs in module_call_specs.items():
1378 module_call_specs=gm_torch_level.meta["module_call_specs"],
1579 module_call_specs: Dict[str, Dict[str, pytree.TreeSpec]] = {}
1609 wrapped_mod, new_preserved_call_signatures, module_call_specs
[all …]
|