Added missing imports to test_statespace_JAX

pymc-devs · Dec 11, 2024 · 1658fe7 · 1658fe7
1 parent 5dd7406
commit 1658fe7
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/tests/statespace/test_statespace_JAX.py b/tests/statespace/test_statespace_JAX.py
@@ -9,10 +9,13 @@
 from pymc.model.transform.optimization import freeze_dims_and_data
 
 from pymc_extras.statespace.utils.constants import (
+    FILTER_OUTPUT_NAMES,
     JITTER_DEFAULT,
     LONG_MATRIX_NAMES,
+    MATRIX_NAMES,
     MISSING_FILL,
     SHORT_NAME_TO_LONG,
+    SMOOTHER_OUTPUT_NAMES,
 )
 from tests.statespace.test_statespace import (  # pylint: disable=unused-import
     exog_ss_mod,