update docstring to clarify behavior

lnccbrown · Jul 12, 2024 · dcc208c · dcc208c
1 parent 3ed3418
commit dcc208c
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/src/hssm/hssm.py b/src/hssm/hssm.py
@@ -460,8 +460,9 @@ def sample(
         Returns
         -------
         az.InferenceData | pm.Approximation
-            A reference to the `model.traces` object, which is an ArviZ `InferenceData`
-            instance if inference_method is `"mcmc"` (default), `"nuts_numpyro"`,
+            A reference to the `model.traces` object, which stores the traces of the
+            last call to `model.sample()`. `model.traces` is an ArviZ `InferenceData`
+            instance if `sampler` is `"mcmc"` (default), `"nuts_numpyro"`,
             `"nuts_blackjax"` or "`laplace"`, or an `Approximation` object if `"vi"`.
         """
         # If initvals are None (default)