Added missing imports to test_statespace_JAX

Author: fonnesbeck

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,