pytorch  1.8.2
About: PyTorch provides Tensor computation (like NumPy) with strong GPU acceleration and Deep Neural Networks (in Python) built on a tape-based autograd system. LTS (Long Term Support) release.
  Fossies Dox: pytorch-1.8.2.tar.gz  ("unofficial" and yet experimental doxygen-generated source code documentation)  

torch.utils._pytree Namespace Reference


class  LeafSpec
class  NodeDef
class  TreeSpec


None _register_pytree_node (Any typ, FlattenFunc flatten_fn, UnflattenFunc unflatten_fn)
Tuple[List[Any], Context_dict_flatten (Dict[Any, Any] d)
Dict[Any, Any] _dict_unflatten (List[Any] values, Context context)
Tuple[List[Any], Context_list_flatten (List[Any] d)
List[Any] _list_unflatten (List[Any] values, Context context)
Tuple[List[Any], Context_tuple_flatten (Tuple[Any,...] d)
Tuple[Any,...] _tuple_unflatten (List[Any] values, Context context)
bool _is_leaf (PyTree pytree)
Tuple[List[Any], TreeSpectree_flatten (PyTree pytree)
PyTree tree_unflatten (List[Any] values, TreeSpec spec)
Optional[List[Any]] _broadcast_to_and_flatten (PyTree pytree, TreeSpec spec)


 Context = Any
 PyTree = Any
 FlattenFunc = Callable[[PyTree], Tuple[List, Context]]
 UnflattenFunc = Callable[[List, Context], PyTree]

Function Documentation

◆ _broadcast_to_and_flatten()

Optional[List[Any]] torch.utils._pytree._broadcast_to_and_flatten ( PyTree  pytree,
TreeSpec  spec 

◆ _dict_flatten()

Tuple[List[Any], Context] torch.utils._pytree._dict_flatten ( Dict[Any, Any]  d)

Definition at line 41 of file

References torch.utils._pytree._dict_flatten(), and c10::prim.list().

Referenced by torch.utils._pytree._dict_flatten().

◆ _dict_unflatten()

Dict[Any, Any] torch.utils._pytree._dict_unflatten ( List[Any]  values,
Context  context 

◆ _is_leaf()

bool torch.utils._pytree._is_leaf ( PyTree  pytree)

◆ _list_flatten()

Tuple[List[Any], Context] torch.utils._pytree._list_flatten ( List[Any]  d)

Definition at line 47 of file

References torch.utils._pytree._list_flatten().

Referenced by torch.utils._pytree._list_flatten().

◆ _list_unflatten()

List[Any] torch.utils._pytree._list_unflatten ( List[Any]  values,
Context  context 

◆ _register_pytree_node()

None torch.utils._pytree._register_pytree_node ( Any  typ,
FlattenFunc  flatten_fn,
UnflattenFunc  unflatten_fn 

◆ _tuple_flatten()

Tuple[List[Any], Context] torch.utils._pytree._tuple_flatten ( Tuple[Any, ...]  d)

◆ _tuple_unflatten()

Tuple[Any, ...] torch.utils._pytree._tuple_unflatten ( List[Any]  values,
Context  context 

◆ tree_flatten()

Tuple[List[Any], TreeSpec] torch.utils._pytree.tree_flatten ( PyTree  pytree)
Flattens a pytree into a list of values and a TreeSpec that can be used
to reconstruct the pytree.

Definition at line 104 of file

References torch.utils._pytree._is_leaf(), torch.utils._pytree.tree_flatten(), and type.

Referenced by torch._vmap_internals._create_batched_inputs(), and torch.utils._pytree.tree_flatten().

◆ tree_unflatten()

PyTree torch.utils._pytree.tree_unflatten ( List[Any]  values,
TreeSpec  spec 
Given a list of values and a TreeSpec, builds a pytree.
This is the inverse operation of `tree_flatten`.

Definition at line 125 of file

References c10::prim.isinstance(), c10::aten.len(), and torch.utils._pytree.tree_unflatten().

Referenced by torch._vmap_internals._create_batched_inputs(), and torch.utils._pytree.tree_unflatten().

Variable Documentation

◆ Context

◆ FlattenFunc

torch.utils._pytree.FlattenFunc = Callable[[PyTree], Tuple[List, Context]]

Definition at line 29 of file

◆ PyTree

torch.utils._pytree.PyTree = Any

Definition at line 28 of file

◆ UnflattenFunc

torch.utils._pytree.UnflattenFunc = Callable[[List, Context], PyTree]

Definition at line 30 of file