diff --git a/src/MEDS_tabular_automl/generate_summarized_reps.py b/src/MEDS_tabular_automl/generate_summarized_reps.py
index 0745440..b0d66ac 100644
--- a/src/MEDS_tabular_automl/generate_summarized_reps.py
+++ b/src/MEDS_tabular_automl/generate_summarized_reps.py
@@ -16,7 +16,7 @@
 )
 
 
-def sparse_aggregate(sparse_matrix: sparray, agg: str) -> csr_array:
+def sparse_aggregate(sparse_matrix: sparray, agg: str) -> np.ndarray | coo_array:
     """Aggregates values in a sparse matrix according to the specified method.
 
     Args: