Fix save_from_both_policies in presence of `save_and_offload_only_t…

…hese_names`

PiperOrigin-RevId: 706720875

jax/_src/ad_checkpoint.py

@@ -142,8 +142,9 @@ def policy(prim, *_, **params):
 def save_from_both_policies(policy_1, policy_2):
 
   def policy(prim, *args, **params):
-    return policy_1(prim, *args, **params) or policy_2(prim, *args, **params)
-
+    out1 = pe.ensure_enum(policy_1(prim, *args, **params))
+    out2 = pe.ensure_enum(policy_2(prim, *args, **params))
+    return isinstance(out1, pe.SaveableType) or isinstance(out2, pe.SaveableType)
   return policy