diff --git a/src/Psl/Arr/group_by.php b/src/Psl/Arr/group_by.php
index e829c09..b5196f5 100644
--- a/src/Psl/Arr/group_by.php
+++ b/src/Psl/Arr/group_by.php
@@ -45,6 +45,8 @@ function group_by(iterable $values, callable $key_func): array
             continue;
         }
 
+        Psl\invariant(is_arraykey($key), 'Expected $key_func to return a value of type array-key, value of type (%s) returned.', gettype($key));
+        /** @psalm-var Tk $key */
         $result[$key] = $result[$key] ?? [];
         $result[$key][] = $value;
     }