diff --git a/docs/source/changes.md b/docs/source/changes.md
index 19698bc5..6674f1a8 100644
--- a/docs/source/changes.md
+++ b/docs/source/changes.md
@@ -12,6 +12,9 @@ releases are available on [PyPI](https://pypi.org/project/pytask) and
   is not present.
 - {pull}`387` replaces pony with sqlalchemy.
 - {pull}`391` removes `@pytask.mark.parametrize`.
+- {pull}`394` allows to add products with {obj}`typing.Annotation` and
+  {obj}`~pytask.Product`.
+- {pull}`395` refactors all occurrences of pybaum to {mod}`_pytask.tree_util`.
 
 ## 0.3.2 - 2023-06-07
 
diff --git a/src/_pytask/clean.py b/src/_pytask/clean.py
index 2004670e..d0aee55d 100644
--- a/src/_pytask/clean.py
+++ b/src/_pytask/clean.py
@@ -29,8 +29,8 @@
 from _pytask.session import Session
 from _pytask.shared import to_list
 from _pytask.traceback import render_exc_info
+from _pytask.tree_util import tree_just_yield
 from attrs import define
-from pybaum.tree_util import tree_just_yield
 
 
 if TYPE_CHECKING:
diff --git a/src/_pytask/collect_command.py b/src/_pytask/collect_command.py
index e297c823..9fda4258 100644
--- a/src/_pytask/collect_command.py
+++ b/src/_pytask/collect_command.py
@@ -26,7 +26,7 @@
 from _pytask.path import relative_to
 from _pytask.pluginmanager import get_plugin_manager
 from _pytask.session import Session
-from pybaum.tree_util import tree_just_flatten
+from _pytask.tree_util import tree_just_flatten
 from rich.text import Text
 from rich.tree import Tree
 
diff --git a/src/_pytask/collect_utils.py b/src/_pytask/collect_utils.py
index 0b179e8d..c82ceb97 100644
--- a/src/_pytask/collect_utils.py
+++ b/src/_pytask/collect_utils.py
@@ -19,9 +19,9 @@
 from _pytask.nodes import PythonNode
 from _pytask.shared import find_duplicates
 from _pytask.task_utils import parse_keyword_arguments_from_signature_defaults
+from _pytask.tree_util import tree_map
 from attrs import define
 from attrs import field
-from pybaum.tree_util import tree_map
 from typing_extensions import Annotated
 from typing_extensions import get_origin
 
diff --git a/src/_pytask/dag.py b/src/_pytask/dag.py
index bbf199a5..b0865b7b 100644
--- a/src/_pytask/dag.py
+++ b/src/_pytask/dag.py
@@ -31,7 +31,7 @@
 from _pytask.shared import reduce_names_of_multiple_nodes
 from _pytask.shared import reduce_node_name
 from _pytask.traceback import render_exc_info
-from pybaum import tree_map
+from _pytask.tree_util import tree_map
 from rich.text import Text
 from rich.tree import Tree
 
diff --git a/src/_pytask/execute.py b/src/_pytask/execute.py
index cce8a054..bcc3fb41 100644
--- a/src/_pytask/execute.py
+++ b/src/_pytask/execute.py
@@ -33,7 +33,7 @@
 from _pytask.traceback import format_exception_without_traceback
 from _pytask.traceback import remove_traceback_from_exc_info
 from _pytask.traceback import render_exc_info
-from pybaum.tree_util import tree_map
+from _pytask.tree_util import tree_map
 from rich.text import Text
 
 
diff --git a/src/_pytask/traceback.py b/src/_pytask/traceback.py
index cc14c181..88c23aef 100644
--- a/src/_pytask/traceback.py
+++ b/src/_pytask/traceback.py
@@ -10,7 +10,7 @@
 
 import _pytask
 import pluggy
-import pybaum
+from _pytask.tree_util import TREE_UTIL_LIB_DIRECTORY
 from rich.traceback import Traceback
 
 
@@ -23,7 +23,6 @@
 
 
 _PLUGGY_DIRECTORY = Path(pluggy.__file__).parent
-_PYBAUM_DIRECTORY = Path(pybaum.__file__).parent
 _PYTASK_DIRECTORY = Path(_pytask.__file__).parent
 
 
@@ -93,7 +92,7 @@ def _is_internal_or_hidden_traceback_frame(
     path = Path(frame.tb_frame.f_code.co_filename)
     return any(
         root in path.parents
-        for root in (_PLUGGY_DIRECTORY, _PYBAUM_DIRECTORY, _PYTASK_DIRECTORY)
+        for root in (_PLUGGY_DIRECTORY, TREE_UTIL_LIB_DIRECTORY, _PYTASK_DIRECTORY)
     )
 
 
diff --git a/src/_pytask/tree_util.py b/src/_pytask/tree_util.py
new file mode 100644
index 00000000..65d17bdd
--- /dev/null
+++ b/src/_pytask/tree_util.py
@@ -0,0 +1,18 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+import pybaum
+from pybaum.tree_util import tree_just_flatten
+from pybaum.tree_util import tree_just_yield
+from pybaum.tree_util import tree_map
+
+
+__all__ = [
+    "tree_just_flatten",
+    "tree_map",
+    "tree_just_yield",
+    "TREE_UTIL_LIB_DIRECTORY",
+]
+
+TREE_UTIL_LIB_DIRECTORY = Path(pybaum.__file__).parent
diff --git a/tests/test_pybaum.py b/tests/test_tree_util.py
similarity index 90%
rename from tests/test_pybaum.py
rename to tests/test_tree_util.py
index b6889fbd..be6212a4 100644
--- a/tests/test_pybaum.py
+++ b/tests/test_tree_util.py
@@ -1,11 +1,11 @@
-"""This module contains tests for pybaum and flexible dependencies and products."""
+"""This module contains tests for tree_util and flexible dependencies and products."""
 from __future__ import annotations
 
 import textwrap
 
 import pytest
 from _pytask.outcomes import ExitCode
-from pybaum import tree_map
+from _pytask.tree_util import tree_map
 from pytask import cli
 from pytask import main
 
@@ -55,11 +55,11 @@ def task_example():
 
 
 @pytest.mark.end_to_end()
-def test_profile_with_pybaum(tmp_path, runner):
+def test_profile_with_pytree(tmp_path, runner):
     source = """
     import time
     import pytask
-    from pybaum.tree_util import tree_just_flatten
+    from _pytask.tree_util import tree_just_flatten
 
     @pytask.mark.produces([{"out_1": "out_1.txt"}, {"out_2": "out_2.txt"}])
     def task_example(produces):