torch.utils._pytree Namespace Reference

Classes
class	LeafSpec
class	NodeDef
class	TreeSpec

Functions
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], TreeSpec]	tree_flatten (PyTree pytree)
PyTree	tree_unflatten (List[Any] values, TreeSpec spec)
Optional[List[Any]]	_broadcast_to_and_flatten (PyTree pytree, TreeSpec spec)

Variables
	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  )

private

Definition at line 163 of file _pytree.py.

References torch.utils._pytree._broadcast_to_and_flatten(), torch.utils._pytree._is_leaf(), c10::prim.isinstance(), c10::aten.len(), type, and c10::prim.zip().

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

_dict_flatten()

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

private

Definition at line 41 of file _pytree.py.

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  )

private

Definition at line 44 of file _pytree.py.

References torch.utils._pytree._dict_unflatten(), and c10::prim.zip().

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

_is_leaf()

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

private

Definition at line 65 of file _pytree.py.

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

Referenced by torch.utils._pytree._broadcast_to_and_flatten(), torch.utils._pytree._is_leaf(), and torch.utils._pytree.tree_flatten().

_list_flatten()

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

private

Definition at line 47 of file _pytree.py.

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  )

private

Definition at line 50 of file _pytree.py.

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

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

_register_pytree_node()

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

private

Definition at line 38 of file _pytree.py.

References torch.utils._pytree._register_pytree_node().

Referenced by torch.utils._pytree._register_pytree_node(), and torch.utils._pytree._tuple_unflatten().

_tuple_flatten()

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

private

Definition at line 53 of file _pytree.py.

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

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

_tuple_unflatten()

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

private

Definition at line 56 of file _pytree.py.

References torch.utils._pytree._register_pytree_node(), and torch.utils._pytree._tuple_unflatten().

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

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 _pytree.py.

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 _pytree.py.

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

torch.utils._pytree.Context = Any

Definition at line 27 of file _pytree.py.

Referenced by org.pytorch.PytorchInstrumentedTests.assetFilePath(), at::native::vulkan::api.context(), at::native::vulkan::ops.copy_(), caffe2::QTensorDeserializer< Context >.Deserialize(), caffe2::python::TensorFeeder< Context >.FeedTensor(), main(), caffe2::OpenCLContextSingleton.OpenCLContextSingleton(), and at::native::vulkan::Conv2dOpContext.run().

FlattenFunc

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

Definition at line 29 of file _pytree.py.

PyTree

torch.utils._pytree.PyTree = Any

Definition at line 28 of file _pytree.py.

UnflattenFunc

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

Definition at line 30 of file _pytree.py.