| /aosp_15_r20/external/pytorch/torch/fx/ |
| H A D | _pytree.py | 27 spec: TreeSpec,
54 def _dict_flatten_spec(d: Dict[Any, Any], spec: TreeSpec) -> List[Any]:
58 def _list_flatten_spec(d: List[Any], spec: TreeSpec) -> List[Any]:
62 def _tuple_flatten_spec(d: Tuple[Any], spec: TreeSpec) -> List[Any]:
66 def _namedtuple_flatten_spec(d: NamedTuple, spec: TreeSpec) -> List[Any]:
70 def _dict_flatten_spec_exact_match(d: Dict[Any, Any], spec: TreeSpec) -> bool:
74 def _list_flatten_spec_exact_match(d: List[Any], spec: TreeSpec) -> bool:
78 def _tuple_flatten_spec_exact_match(d: Tuple[Any], spec: TreeSpec) -> bool:
82 def _namedtuple_flatten_spec_exact_match(d: NamedTuple, spec: TreeSpec) -> bool:
|
| /aosp_15_r20/external/pytorch/torch/export/ |
| H A D | _unlift.py | 124 in_spec: pytree.TreeSpec,
125 out_spec: Optional[pytree.TreeSpec],
160 in_spec: pytree.TreeSpec,
161 out_spec: Optional[pytree.TreeSpec],
|
| H A D | _tree_utils.py | 6 def reorder_kwargs(user_kwargs: Dict[str, Any], spec: TreeSpec) -> Dict[str, Any]:
34 spec1: TreeSpec,
35 spec2: TreeSpec,
|
| H A D | _trace.py | 505 in_spec: TreeSpec,
506 out_spec: TreeSpec,
1064 original_in_spec: TreeSpec,
1219 orig_in_spec: TreeSpec,
1247 orig_in_spec: TreeSpec,
1567 orig_in_spec: TreeSpec,
|
| H A D | unflatten.py | 182 target_spec: pytree.TreeSpec,
183 input_spec: pytree.TreeSpec,
|
| /aosp_15_r20/external/pytorch/torch/onnx/_internal/ |
| H A D | io_adapter.py | 139 def _replace_tuple_with_list(spec: pytree.TreeSpec) -> pytree.TreeSpec:
146 def _open_top_level_list_if_single_element(spec: pytree.TreeSpec) -> pytree.TreeSpec:
153 spec1: pytree.TreeSpec, spec2: pytree.TreeSpec, error_message: str
|
| /aosp_15_r20/external/pytorch/torch/utils/ |
| H A D | _pytree.py | 684 class TreeSpec: class
878 def tree_unflatten(leaves: Iterable[Any], treespec: TreeSpec) -> PyTree:
1296 treespec: TreeSpec,
1351 def _treespec_to_json(treespec: TreeSpec) -> _TreeSpecSchema:
1425 def treespec_dumps(treespec: TreeSpec, protocol: Optional[int] = None) -> str:
1463 def treespec_pprint(treespec: TreeSpec) -> str:
1476 def pytree_to_str(treespec: TreeSpec) -> str:
|
| H A D | _cxx_pytree.py | 82 TreeSpec = PyTreeSpec variable
294 def tree_unflatten(leaves: Iterable[Any], treespec: TreeSpec) -> PyTree:
856 treespec: TreeSpec,
867 def treespec_dumps(treespec: TreeSpec, protocol: Optional[int] = None) -> str:
901 def treespec_pprint(treespec: TreeSpec) -> str:
|
| /aosp_15_r20/external/executorch/exir/ |
| H A D | tracer.py | 484 in_spec: Optional[TreeSpec] = None,
496 in_spec: Optional[TreeSpec],
576 in_spec: Optional[TreeSpec] = None,
|
| /aosp_15_r20/external/pytorch/torch/_functorch/_aot_autograd/ |
| H A D | input_output_analysis.py | 421 in_spec: pytree.TreeSpec,
422 out_spec: pytree.TreeSpec,
|
| H A D | schemas.py | 697 in_spec: pytree.TreeSpec,
698 out_spec: pytree.TreeSpec,
|
| H A D | utils.py | 146 def set(self, spec: pytree.TreeSpec) -> None:
|
| /aosp_15_r20/external/executorch/extension/pytree/ |
| H A D | __init__.py | 33 TreeSpec as TreeSpec, unknown
|
| H A D | pytree.h | 285 using TreeSpec = ContainerHandle<TreeSpecLeaf, Aux>; variable
|
| /aosp_15_r20/external/pytorch/test/export/ |
| H A D | test_unflatten.py | 270 target_spec: TreeSpec,
271 input_spec: TreeSpec,
|
| /aosp_15_r20/external/executorch/exir/emit/ |
| H A D | _emitter.py | 1479 def create_container_str(spec: Optional[pytree.TreeSpec]) -> str:
|
| /aosp_15_r20/external/pytorch/torch/_subclasses/ |
| H A D | fake_tensor.py | 2058 args_spec: TreeSpec,
|
| /aosp_15_r20/external/pytorch/torch/_inductor/ |
| H A D | pattern_matcher.py | 572 def norm_spec(s: pytree.TreeSpec) -> pytree.TreeSpec:
|