add typing to fuse

scverse · Dec 16, 2024 · eaf3014 · eaf3014
1 parent 5c7061e
commit eaf3014
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/rapids_singlecell/preprocessing/_harmony/_fuses.py b/src/rapids_singlecell/preprocessing/_harmony/_fuses.py
@@ -4,7 +4,7 @@
 
 
 @cp.fuse
-def _get_factor(O_k, ridge_lambda):
+def _get_factor(O_k: cp.ndarray, ridge_lambda: float) -> cp.ndarray:
     return 1 / (O_k + ridge_lambda)
 
 
@@ -29,5 +29,5 @@ def _log_div_OE(O: cp.ndarray, E: cp.ndarray) -> cp.ndarray:
 
 
 @cp.fuse
-def _R_multi_m(R, other):
+def _R_multi_m(R: cp.ndarray, other: cp.ndarray) -> cp.ndarray:
     return R * 2 * (1 - other)