Home
last modified time | relevance | path

Searched defs:PyTree (Results 1 – 10 of 10) sorted by relevance

/aosp_15_r20/external/pytorch/torch/utils/
H A D_cxx_pytree.py81 PyTree = Any variable
247 tree: PyTree,
248 is_leaf: Optional[Callable[[PyTree], bool]] = None,
322 tree: PyTree,
323 is_leaf: Optional[Callable[[PyTree], bool]] = None,
357 tree: PyTree,
358 is_leaf: Optional[Callable[[PyTree], bool]] = None,
392 tree: PyTree,
393 is_leaf: Optional[Callable[[PyTree], bool]] = None,
428 tree: PyTree,
[all …]
H A D_pytree.py114 PyTree = Any variable
672 def _is_leaf(tree: PyTree, is_leaf: Optional[Callable[[PyTree], bool]] = None) -> bool:
720 def _flatten_up_to_helper(self, tree: PyTree, subtrees: List[PyTree]) -> None:
796 def flatten_up_to(self, tree: PyTree) -> List[PyTree]:
846 tree: PyTree,
848 is_leaf: Optional[Callable[[PyTree], bool]] = None,
867 tree: PyTree,
868 is_leaf: Optional[Callable[[PyTree], bool]] = None,
891 tree: PyTree,
892 is_leaf: Optional[Callable[[PyTree], bool]] = None,
[all …]
/aosp_15_r20/external/pytorch/torch/_higher_order_ops/
H A Dassociative_scan.py84 combine_fn: Callable[[pytree.PyTree, pytree.PyTree], pytree.PyTree],
85 input: pytree.PyTree,
/aosp_15_r20/external/executorch/extension/pytree/
H A Dpybindings.cpp91 class PyTree { class
248 explicit PyTree(PyTreeSpec spec) : spec_(std::move(spec)) {} in PyTree() function in executorch::extension::pytree::__anon59fa55ac0111::PyTree
/aosp_15_r20/external/executorch/exir/passes/
H A Dremove_mixed_type_operators.py64 def try_coerce(value: PyTree, arg: torch.Argument) -> PyTree:
H A Dsym_to_tensor_pass.py51 def try_coerce(value: PyTree, arg: torch.Argument) -> PyTree:
/aosp_15_r20/external/executorch/exir/
H A Dcontrol_flow.py132 init_val: pytree.PyTree,
H A Dpass_base.py725 def update(key: K, args: MutableMapping[K, PyTree], schema: ArgSchema) -> None: argument
/aosp_15_r20/external/pytorch/torch/fx/
H A D_pytree.py26 pytree: PyTree,
/aosp_15_r20/external/pytorch/torch/_subclasses/
H A Dfake_tensor.py236 def tree_flatten_only(ty: Type[T], tree: PyTree) -> List[T]:
2256 args_spec: PyTree,