Type Model.register_rv for for downstream typing

In pymc/distributions/distribution.py, this change allows the type checker to infer that `rv_out` can only be a TensorVariable.
pymc-devs · Sep 4, 2024 · 383d08e · 383d08e
1 parent cb9cc57
commit 383d08e
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/pymc/model/core.py b/pymc/model/core.py
@@ -1166,16 +1166,16 @@ def set_data(
 
     def register_rv(
         self,
-        rv_var,
-        name,
+        rv_var: TensorVariable,
+        name: str,
         *,
         observed=None,
         total_size=None,
         dims=None,
         default_transform=UNSET,
         transform=UNSET,
         initval=None,
-    ):
+    ) -> TensorVariable:
         """Register an (un)observed random variable with the model.
 
         Parameters