Update fixes for b/498090408 and b/498080408 to be more concise.

Author: QuinnMMcGarry

axlearn/common/flash_attention/tpu_paged_attention.py

@@ -258,8 +258,8 @@ def __call__(
 
         in_specs = [
             q_block_spec,  # Query
-            pl.BlockSpec(memory_space=None),  # Key pages
-            pl.BlockSpec(memory_space=None),  # Value pages
+            pl.BlockSpec(memory_space=pl.ANY),  # Key pages
+            pl.BlockSpec(memory_space=pl.ANY),  # Value pages
             bias_spec,  # Bias
             logit_sink_spec,  # Logit sink
         ]

axlearn/common/trainer_test.py

@@ -13,6 +13,7 @@
 import tempfile
 import unittest
 from collections.abc import Sequence
+from packaging import version
 from typing import Any, Callable, Literal, Optional
 
 import chex
@@ -535,11 +536,11 @@ def mock_compile_train_step(*args, compiler_options=None, **kwargs):
                 if not enable_python_cache:
                     # As of Jax >= 0.6.0, enable_python_cache=False no longer affects the
                     # AOT compilation path. We now expect the cache hits to be 0
-                    # if jax.__version__ >= "0.6.0":
-                    pytest.skip(
-                        # pylint: disable-next=line-too-long
-                        "AOT compilation path is not affected by 'enable_python_cache' with Jax >= 0.6.0"
-                    )
+                    if version.parse(jax.__version__) >= version.parse("0.6.0"):
+                        pytest.skip(
+                            # pylint: disable-next=line-too-long
+                            "AOT compilation path is not affected by 'enable_python_cache' with Jax >= 0.6.0"
+                        )
                     # We expect to have hit the lowering cache on all but one step.
                     self.assertEqual(end_cache_hits - start_cache_hits, cfg.max_step - 1)
                     self.assertEqual(mocked_compile_fn.call_count, cfg.max_step)
@@ -551,11 +552,11 @@ def mock_compile_train_step(*args, compiler_options=None, **kwargs):
                 self.assertEqual(compiled_with_options_call_count[0], 2)
             else:
                 if not enable_python_cache:
-                    # if jax.__version__ >= "0.6.0":
-                    pytest.skip(
-                        # pylint: disable-next=line-too-long
-                        "AOT compilation path is not affected by 'enable_python_cache' with Jax >= 0.6.0"
-                    )
+                    if version.parse(jax.__version__) >= version.parse("0.6.0"):
+                        pytest.skip(
+                            # pylint: disable-next=line-too-long
+                            "AOT compilation path is not affected by 'enable_python_cache' with Jax >= 0.6.0"
+                        )
                     self.assertEqual(end_cache_hits - start_cache_hits, cfg.max_step - 1)
                     self.assertEqual(mocked_compile_fn.call_count, cfg.max_step)
                 else:
@@ -607,15 +608,12 @@ def test_compile_train_step(self, *, platform, mesh_shape):
         # In a single-host environment, both compiled functions should match.
         # Skip this part of the test if Jax >= 0.8.2 as the behavior changes
         # TODO(samuel-andersen): Investigate why this behavior has changed
-        # TEMP FOR 0.10.0DEV
-        '''
-        if jax.__version__ < "0.8.2":
+        if version.parse(jax.__version__) < version.parse("0.8.2"):
             self.assertEqual(compiled_without_args.as_text(), compiled_with_input_batch.as_text())
             self.assertEqual(
                 aot_model_analysis(compiled_without_args),
                 aot_model_analysis(compiled_with_input_batch),
             )
-        '''
 
         # A version compiled with non-default compiled args should be different.
         compiled_with_compiler_options = trainer.compile_train_step(
@@ -628,9 +626,7 @@ def test_compile_train_step(self, *, platform, mesh_shape):
             trainer_state=trainer.trainer_state, input_batch=input_batch
         )
         # Skip this part of the test if Jax >= 0.8.2 as the behavior changes
-        # TEMP FOR 0.10.0DEV
-        '''
-        if jax.__version__ < "0.8.2":
+        if version.parse(jax.__version__) < version.parse("0.8.2"):
             self.assertEqual(
                 compiled_without_args.as_text(),
                 compiled_with_trainer_state_and_input_batch.as_text(),
@@ -639,7 +635,7 @@ def test_compile_train_step(self, *, platform, mesh_shape):
                 aot_model_analysis(compiled_without_args),
                 aot_model_analysis(compiled_with_trainer_state_and_input_batch),
             )
-        '''
+
     @parameterized.parameters(
         {"return_evaler_summaries": None},
         {"return_evaler_summaries": True},