Using version guard to fix change that broke backwards compatibility …

…with some older versions of JAX.

PiperOrigin-RevId: 694248080

kfac_jax/_src/utils/parallel.py

@@ -22,6 +22,11 @@
 
 from kfac_jax._src.utils import types
 
+jax_version = (
+    jax.__version_info__ if hasattr(jax, "__version_info__")
+    else tuple(map(int, jax.__version__.split("."))))
+
+
 Array = types.Array
 Numeric = types.Numeric
 PRNGKey = types.PRNGKey
@@ -35,7 +40,19 @@ def in_pmap(axis_name: str | None) -> bool:
   if axis_name is None:
     return False
 
-  return axis_name in core.unsafe_get_axis_names_DO_NOT_USE()
+  if jax_version >= (0, 4, 36):
+    return axis_name in core.unsafe_get_axis_names_DO_NOT_USE()
+
+  else:
+    try:
+      # The only way to know if we are under `jax.pmap` is to check if the
+      # function call below raises a `NameError` or not.
+      core.axis_frame(axis_name)
+
+      return True
+
+    except NameError:
+      return False
 
 
 def wrap_if_pmap(