[RLlib] Fix JAX import bug. (#12621)

rllib/utils/framework.py

@@ -281,7 +281,7 @@ def get_activation_fn(name: Optional[str] = None, framework: str = "tf"):
     elif framework == "jax":
         if name in ["linear", None]:
             return None
-        jax = try_import_jax()
+        jax, flax = try_import_jax()
         if name == "swish":
             return jax.nn.swish
         if name == "relu":

rllib/utils/test_utils.py

@@ -5,7 +5,7 @@ import numpy as np
 from ray.rllib.utils.framework import try_import_jax, try_import_tf, \
     try_import_torch
 
-jax = try_import_jax()
+jax, flax = try_import_jax()
 tf1, tf, tfv = try_import_tf()
 if tf1:
     eager_mode = None