deepfold.utils.tensor_utils.tree_map¶

deepfold.utils.tensor_utils.tree_map(fn: Callable[[T], Any], tree: T, leaf_type: Type[T]) → Any[source]¶
deepfold.utils.tensor_utils.tree_map(fn: Callable[[T], Any], tree: dict, leaf_type: Type[T]) → dict
deepfold.utils.tensor_utils.tree_map(fn: Callable[[T], Any], tree: list, leaf_type: Type[T]) → list
deepfold.utils.tensor_utils.tree_map(fn: Callable[[T], Any], tree: tuple, leaf_type: Type[T]) → tuple