diff --git a/pytools/__init__.py b/pytools/__init__.py
index 2694de21d112dfd5db8989165029f4d3b93634e0..1b1608b0dc445141a03807a913ec2230367823ee 100644
--- a/pytools/__init__.py
+++ b/pytools/__init__.py
@@ -950,7 +950,7 @@ class memoize_in:  # noqa: N801
             try:
                 return self.cache_dict[args]
             except KeyError:
-                result = inner(*args)
+                result = inner(*args, **kwargs)
                 self.cache_dict[args] = result
                 return result
 
@@ -984,12 +984,12 @@ class keyed_memoize_in(Generic[P, R]):  # noqa: N801
         @wraps(inner)
         def new_inner(*args: P.args, **kwargs: P.kwargs) -> R:
             assert not kwargs
-            key = self.key(*args)
+            key = self.key(*args, **kwargs)
 
             try:
                 return self.cache_dict[key]
             except KeyError:
-                result = inner(*args)
+                result = inner(*args, **kwargs)
                 self.cache_dict[key] = result
                 return result
 
diff --git a/pytools/graph.py b/pytools/graph.py
index f00547beb5c3f4f40eaa0b631ce626246ca8b8af..359e8060e3b0e594b82e5886102ede126d01573e 100644
--- a/pytools/graph.py
+++ b/pytools/graph.py
@@ -258,7 +258,7 @@ class CycleError(Exception):
 
 
 class _SupportsLT(Protocol):
-    def __lt__(self, other: object) -> bool:
+    def __lt__(self, other: Any) -> bool:
         ...
 
 
@@ -278,9 +278,10 @@ class _HeapEntry(Generic[NodeT]):
         return self.key < other.key
 
 
-def compute_topological_order(graph: GraphT[NodeT],
-                              key: Optional[Callable[[NodeT], Any]] = None) \
-                              -> List[NodeT]:
+def compute_topological_order(
+        graph: GraphT[NodeT],
+        key: Optional[Callable[[NodeT], _SupportsLT]] = None,
+        ) -> List[NodeT]:
     """Compute a topological order of nodes in a directed graph.
 
     :arg key: A custom key function may be supplied to determine the order in